• Search by Module
• Search by Words
• Search Projects

• Java
• Python
• JavaScript
• TypeScript
• C++
• Scala
• Blog

• tensor2tensor-master
  • AUTHORS
  • floyd.yml
  • tensor2tensor
    • serving
      • export.py
      • query.py
      • __init__.py
      • README.md
      • serving_utils.py
    • utils
      • video2gif.py
      • mtf_model.py
      • multistep_with_adamoptimizer_test.py
      • avg_checkpoints.py
      • expert_utils_test.py
      • data_reader_test.py
      • test_utils.py
      • quantization.py
      • trainer_lib_test.py
      • mlperf_log.py
      • hparams_lib_test.py
      • hparam_test.py
      • usr_dir.py
      • get_rouge.py
      • sari_hook.py
      • metrics_test.py
      • learning_rate.py
      • devices.py
      • yellowfin_test.py
      • misc_utils_test.py
      • metrics.py
      • scheduled_sampling.py
      • cloud_mlengine.py
      • decoding.py
      • adv_attack_utils.py
      • sari_hook_test.py
      • test_utils_test.py
      • rouge.py
      • video_metrics.py
      • diet_test.py
      • flags.py
      • adafactor.py
      • get_ende_bleu.sh
      • diet.py
      • data_reader.py
      • rouge_test.py
      • registry.py
      • restore_hook.py
      • checkpoint_compatibility_test.py
      • hparams_lib.py
      • multistep_with_adamoptimizer.py
      • t2t_model.py
      • yellowfin.py
      • __init__.py
      • multistep_optimizer_test.py
      • registry_test.py
      • misc_utils.py
      • video_metrics_test.py
      • bleu_hook_test.py
      • mlperf_tags.py
      • optimize.py
      • partial_checkpoint_load_hook.py
      • hparam.py
      • expert_utils.py
      • pruning_utils.py
      • compute_video_metrics.py
      • optimize_test.py
      • t2t_model_test.py
      • metrics_hook_test.py
      • beam_search.py
      • contrib.py
      • beam_search_test.py
      • trainer_lib.py
      • update_ops_hook.py
      • get_cnndm_rouge.sh
      • metrics_hook.py
      • multistep_optimizer.py
      • video
        • reward_confusion.py
        • prediction2gif.py
      • bleu_hook.py
    • layers
      • common_video.py
      • discretization.py
      • latent_layers.py
      • transformer_glow_layers_ops_test.py
      • transformer_glow_layers.py
      • transformer_memory.py
      • discretization_test.py
      • transformer_glow_layers_test.py
      • transformer_memory_test.py
      • common_hparams.py
      • modalities_test.py
      • area_attention_test.py
      • message_passing_attention.py
      • latent_layers_test.py
      • transformer_glow_layers_ops.py
      • common_image_attention.py
      • common_layers.py
      • ngram.py
      • common_attention_test.py
      • common_image_attention_test.py
      • ngram_test.py
      • vqa_layers.py
      • common_video_test.py
      • __init__.py
      • vq_discrete.py
      • modalities.py
      • common_layers_test.py
      • area_attention.py
      • transformer_layers.py
      • common_audio.py
    • problems_colab.py
    • models
      • revnet_test.py
      • slicenet.py
      • mtf_image_transformer.py
      • distillation.py
      • neural_assistant.py
      • transformer.py
      • basic_test.py
      • slicenet_test.py
      • neural_gpu.py
      • basic.py
      • bytenet_test.py
      • image_transformer_2d.py
      • mtf_resnet.py
      • text_cnn.py
      • evolved_transformer.py
      • xception.py
      • neural_architecture_search
        • nas_model_test.py
        • nas_layers.py
        • __init__.py
        • README.md
        • nas_model.py
        • nas_layers_test.py
      • evolved_transformer_test.py
      • mtf_transformer_test.py
      • bytenet.py
      • image_transformer_2d_test.py
      • resnet.py
      • lstm_test.py
      • image_transformer.py
      • lstm.py
      • mtf_transformer2.py
      • mtf_image_transformer_test.py
      • xception_test.py
      • __init__.py
      • README.md
      • shake_shake.py
      • transformer_test.py
      • resnet_test.py
      • research
        • universal_transformer_util.py
        • rl.py
        • transformer_vae_flow_prior.py
        • moe_experiments.py
        • autoencoders.py
        • similarity_transformer.py
        • glow_test.py
        • transformer_moe.py
        • neural_stack_test.py
        • vqa_attention.py
        • transformer_aux.py
        • cycle_gan.py
        • multiquery_paper.py
        • autoencoders_test.py
        • gene_expression.py
        • glow_ops_test.py
        • attention_lm_moe.py
        • shuffle_network.py
        • transformer_vae_test.py
        • gene_expression_test.py
        • transformer_revnet.py
        • adafactor_experiments.py
        • transformer_nat.py
        • vqa_attention_test.py
        • glow_ops.py
        • transformer_revnet_test.py
        • transformer_vae.py
        • attention_lm.py
        • universal_transformer_test.py
        • moe.py
        • transformer_vae_flow_prior_ops.py
        • lm_experiments.py
        • transformer_symshard.py
        • super_lm.py
        • glow.py
        • glow_init_hook.py
        • transformer_sketch.py
        • __init__.py
        • residual_shuffle_exchange.py
        • universal_transformer.py
        • vqa_recurrent_self_attention.py
        • neural_stack.py
        • vqa_self_attention.py
        • aligned.py
        • transformer_parallel.py
        • transformer_aux_test.py
      • neural_gpu_test.py
      • revnet.py
      • image_transformer_test.py
      • vanilla_gan.py
      • mtf_transformer.py
      • video
        • nfg_conv3d_test.py
        • sv2p_test.py
        • savp.py
        • epva_params.py
        • emily.py
        • base_vae.py
        • nfg_conv_test.py
        • nfg_uncond_test.py
        • next_frame_glow.py
        • sv2p_params.py
        • basic_deterministic_test.py
        • emily_test.py
        • nfg_test_utils.py
        • basic_stochastic.py
        • basic_deterministic_params.py
        • basic_recurrent.py
        • sv2p.py
        • savp_params.py
        • nfg_conv_lstm_test.py
        • basic_recurrent_test.py
        • epva.py
        • savp_test.py
        • basic_stochastic_test.py
        • __init__.py
        • nfg_interpolate.py
        • basic_deterministic.py
        • base.py
        • tests_utils.py
    • rl
      • gym_utils_test.py
      • gym_utils.py
      • evaluator.py
      • trainer_model_free_tictactoe_test.py
      • evaluator_test.py
      • batch_runner_test.py
      • dopamine_connector.py
      • ppo.py
      • trainer_model_free.py
      • restarter.py
      • restarter_test.py
      • rl_utils.py
      • trainer_model_based.py
      • trainer_model_based_agent_only.py
      • batch_dqn_agent_test.py
      • trainer_model_based_params.py
      • trainer_model_free_test.py
      • player.py
      • trainer_model_based_recurrent_test.py
      • __init__.py
      • README.md
      • datagen_with_agent.py
      • player_utils.py
      • ppo_learner.py
      • trainer_model_based_sv2p_test.py
      • trainer_model_based_stochastic_test.py
      • policy_learner.py
      • trainer_model_based_test.py
      • envs
        • py_func_batch_env.py
        • simulated_batch_gym_env.py
        • tf_atari_wrappers.py
        • in_graph_batch_env.py
        • __init__.py
        • simulated_batch_env.py
    • test_data
      • example_usr_dir
        • my_submodule.py
        • __init__.py
        • requirements.txt
      • vocab.translate_ende_wmt8k.8192.subwords
      • transformer_test_ckpt
        • flags.txt
        • model.ckpt-1.index
        • hparams.json
        • checkpoint
        • model.ckpt-1.data-00000-of-00002
    • notebooks
      • asr_transformer.ipynb
      • t2t_problem.ipynb
      • Transformer_translate.ipynb
    • data_generators
      • wiki_lm.py
      • translate_test.py
      • translate_enmk.py
      • audio_encoder.py
      • snli.py
      • translate_enzh.py
      • mscoco_test.py
      • speech_recognition.py
      • conll_ner.py
      • audio_test.py
      • imdb.py
      • gym_env_test.py
      • algorithmic_math_deepmind.py
      • moving_mnist.py
      • timeseries_data_generator.py
      • desc2code.py
      • dialog_opensubtitles.py
      • function_docstring.py
      • wsj_parsing.py
      • tokenizer.py
      • cipher.py
      • gene_expression.py
      • allen_brain.py
      • wiki_multi_problems.py
      • wikitext103.py
      • mnist.py
      • celeba.py
      • librispeech.py
      • translate_ende.py
      • timeseries.py
      • translate_entn.py
      • gene_expression_test.py
      • ptb.py
      • problem_hparams.py
      • paraphrase_ms_coco_test.py
      • dna_encoder_test.py
      • dialog_abstract.py
      • translate_enid.py
      • text_encoder_test.py
      • translate_enes.py
      • lm1b.py
      • generator_utils.py
      • dna_encoder.py
      • cnn_dailymail.py
      • wnli.py
      • transduction_problems.py
      • style_transfer_test.py
      • fsns.py
      • mscoco.py
      • celeba_test.py
      • imagenet_test.py
      • babi_qa.py
      • desc2code_test.py
      • video_utils_test.py
      • vqa.py
      • text_encoder_build_subword.py
      • wiki_revision_utils.py
      • algorithmic_test.py
      • multi_problem_v2.py
      • tokenizer_test.py
      • video_generated.py
      • program_search.py
      • image_utils.py
      • mrpc.py
      • transduction_problems_test.py
      • multi_problem_v2_test.py
      • translate_enro.py
      • audio.py
      • ice_parsing.py
      • image_lsun.py
      • cleaner_en_xx.py
      • algorithmic_math.py
      • scitail.py
      • style_transfer.py
      • test_data
        • vocab-1.txt
        • vocab-2.txt
        • 1.csv
        • corpus-1.txt
        • corpus-2.txt
      • multi_problem.py
      • algorithmic.py
      • rte.py
      • allen_brain_test.py
      • lm1b_mnli.py
      • video_utils.py
      • dialog_personachat.py
      • wikifact
        • README.md
      • paraphrase_ms_coco.py
      • ocr.py
      • problem.py
      • dialog_cornell.py
      • wiki_revision.py
      • stanford_nli.py
      • subject_verb_agreement.py
      • gym_env.py
      • timeseries_data_generator_test.py
      • bair_robot_pushing.py
      • sst_binary.py
      • image_utils_test.py
      • text_problems.py
      • pointer_generator_word.py
      • translate_enet.py
      • generator_utils_test.py
      • __init__.py
      • text_problems_test.py
      • common_voice.py
      • README.md
      • google_robot_pushing.py
      • translate_ende_test.py
      • yelp_polarity.py
      • timeseries_test.py
      • text_encoder.py
      • program_search_test.py
      • ops
        • subword_text_encoder.cc
        • subword_text_encoder_ops.cc
        • testdata
          • subwords
        • subword_text_encoder_test.cc
        • subword_text_encoder.h
        • pack_sequences_ops.cc
        • pack_sequences_ops_test.py
        • subword_text_encoder_ops_test.py
      • wikisum
        • delete_instances.sh
        • html.py
        • validate_data.py
        • wikisum.py
        • utils_test.py
        • parallel_launch.py
        • get_references_web_single_group.py
        • test_data
          • para_good1.txt
          • para_bad1.txt
        • get_references_commoncrawl.py
        • get_references_web.py
        • produce_examples.py
        • __init__.py
        • README.md
        • utils.py
        • generate_vocab.py
      • translate_envi.py
      • multinli.py
      • common_voice_test.py
      • problem_test.py
      • all_problems.py
      • algorithmic_math_test.py
      • celebahq.py
      • vqa_utils.py
      • translate.py
      • lm1b_imdb.py
      • dialog_dailydialog.py
      • inspect_tfrecord.py
      • cola.py
      • enwik8.py
      • imagenet.py
      • cifar.py
      • yelp_full.py
      • algorithmic_math_two_variables.py
      • lambada.py
      • wiki.py
      • qnli.py
      • translate_enfr.py
      • quora_qpairs.py
      • translate_encs.py
      • squad.py
    • __init__.py
    • visualization
      • visualization_test.py
      • attention.py
      • TransformerVisualization.ipynb
      • attention.js
      • __init__.py
      • visualization.py
    • bin
      • t2t-eval
      • t2t-trainer
      • t2t_trainer.py
      • t2t_translate_all.py
      • t2t-translate-all
      • t2t_prune.py
      • t2t_datagen.py
      • t2t-decoder
      • t2t-make-tf-configs
      • t2t-avg-all
      • build_vocab.py
      • make_tf_configs.py
      • t2t_avg_all.py
      • t2t_trainer_test.py
      • t2t_distill.py
      • t2t_attack.py
      • t2t-exporter
      • t2t-bleu
      • __init__.py
      • t2t_eval.py
      • t2t-insights-server
      • t2t_bleu.py
      • t2t_decoder.py
      • t2t-datagen
      • t2t-query-server
    • metrics
      • video_conditional_fvd_test.py
      • video_conditional_fvd.py
      • __init__.py
    • problems_test.py
    • problems.py
    • insights
      • query_processor.py
      • graph.py
      • insight_configuration.proto
      • server.py
      • polymer
        • explore_view
          • explore-view.js
          • explore-view.html
        • query_card
          • query-card.js
          • query-card.html
        • insights_app
          • insights-app.html
          • insights-app.js
        • language_selector
          • language-selector-content.html
          • language-selector-content.js
          • language-selector.html
          • language-selector.js
        • attention_visualization
          • attention-visualization.js
          • attention-visualization.html
        • processing_visualization
          • processing-visualization.js
          • processing-visualization.html
        • bower.json
        • translation_result
          • translation-result.html
          • translation-result.js
        • common-types.js
        • graph_visualization
          • graph-visualization.html
          • graph-visualization.js
        • .bowerrc
        • tensor2tensor.html
        • index.html
      • __init__.py
      • README.md
      • transformer_model.py
    • envs
      • time_step_test.py
      • mujoco_problems.py
      • env_problem_utils.py
      • trajectory.py
      • gym_env_problem.py
      • gym_env_problem_test.py
      • tic_tac_toe_env_test.py
      • tic_tac_toe_env_problem.py
      • time_step.py
      • trajectory_test.py
      • env_problem.py
      • tic_tac_toe_env.py
      • gym_spaces_utils_test.py
      • env_problem_utils_test.py
      • tic_tac_toe_env_problem_test.py
      • __init__.py
      • mujoco_problems_test.py
      • rendered_env_problem_test.py
      • rendered_env_problem.py
      • gym_spaces_utils.py
  • LICENSE
  • CONTRIBUTING.md
  • ISSUE_TEMPLATE.md
  • setup.py
  • .travis.yml
  • README.md
  • floyd_requirements.txt
  • pylintrc
  • .gitignore
  • docs
    • multi_problem.md
    • new_model.md
    • walkthrough.md
    • cloud_mlengine.md
    • index.md
    • overview.md
    • tutorials
      • asr_with_transformer.md
    • distributed_training.md
    • new_problem.md
    • cloud_tpu.md
  • oss_scripts
    • oss_integration_test.sh
    • oss_pip_install.sh
    • oss_tests.sh
    • oss_release.sh

# 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 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.compat.v1 as tf

BATCH_SIZE = 3
INPUT_LENGTH = 5
TARGET_LENGTH = 7
VOCAB_SIZE = 10
DECODE_LENGTH = 3


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 = num_decoder_layers

  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))
  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),
  }
  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)


class EvolvedTransformerTest(tf.test.TestCase):

  def testEvolvedTransformer(self):
    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):
    tf.set_random_seed(1234)
    model, features = get_model(transformer.transformer_tiny())

    decode_length = DECODE_LENGTH

    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 = DECODE_LENGTH

    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 = DECODE_LENGTH

    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 = 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(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):
    tf.set_random_seed(1234)
    decode_length = 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(
          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):
    tf.set_random_seed(1234)
    decode_length = 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(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)

  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()