text: # coding=utf-8
  # Copyright 2018 The Tensor2Tensor Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT 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
  decoder: [Self-Attention, Source-Target-Attention, Feed-forward] x n
  
  """
  
  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_image_attention as cia
  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 Imagetransformer(t2t_model.T2TModel):
    """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)
      targets = features["targets"]
      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.
      losses = []
      # Prepare decoder inputs and bias.
      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])
      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,
          losses=losses,
          name="decoder")
      output = cia.create_output(decoder_output, rows, cols, targets, hparams)
  
      if losses:
        return output, {"extra_loss": tf.add_n(losses)}
      else:
        return output
  
    def loss(self, logits, features):
      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.
  
      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`}.
      """
      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)
  
    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
  class ImagetransformerMoe(t2t_model.T2TModel):
    """Conditional image generation with attention and MoE."""
  
    @property
    def use_body_sharded(self):
      return True
  
    def body_sharded(self, sharded_features):
      dp = self._data_parallelism
      hparams = copy.copy(self._hparams)
      inputs = sharded_features["inputs"]
      targets = sharded_features["targets"]
  
      # Determine attention type and padding from hparams.
      q_padding, kv_padding = "VALID", "VALID"
      if hparams.q_filter_width > 1:
        q_padding = "LEFT"
      if hparams.kv_filter_width > 1:
        kv_padding = "LEFT"
  
      # Prepare decoder inputs and bias.
      decoder_input, rows, cols = dp(cia.prepare_decoder_inputs,
                                     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,
          decoder_input,
          hparams.num_hidden_layers,
          hparams,
          self_attention_bias=None,
          enc_output=None,
          attention_type=hparams.dec_attention_type,
          name="decoder")
  
      output = dp(cia.create_output, decoder_output, rows, cols, targets, hparams)
      return output, extra_loss
  
  
  @registry.register_hparams
  def image_transformer_base():
    """Set of hyperparameters."""
    hparams = common_hparams.basic_params1()
    hparams.hidden_size = 512
    hparams.batch_size = 4
    hparams.max_length = 3075
    hparams.dropout = 0.0
    hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
    hparams.optimizer_adam_epsilon = 1e-9
    hparams.learning_rate_decay_scheme = "noam"
    hparams.learning_rate = 0.1
    hparams.learning_rate_warmup_steps = 4000
    hparams.initializer_gain = 0.2
    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.label_smoothing = 0.0
    hparams.target_modality = "image:identity"
    hparams.norm_type = "layer"
    hparams.layer_prepostprocess_dropout = 0.0
    hparams.add_hparam("filter_size", 512)  # Add new ones like this.
  
    # 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("ffn_layer", "conv_hidden_relu")
    # All hyperparameters ending in "dropout" are automatically set to 0.0
    # when not in training mode.
    hparams.add_hparam("attention_dropout", 0.0)
    hparams.add_hparam("relu_dropout", 0.0)
    hparams.add_hparam("pos", "timing")  # timing, none
    hparams.add_hparam("nbr_decoder_problems", 1)
    hparams.add_hparam("num_output_layers", 3)
    hparams.add_hparam("block_size", 1)
  
    # dilated attention based flags
    hparams.add_hparam("gap_sizes", [2, 4, 8, 16, 32, 64, 2, 4, 8, 16, 32, 64])
  
    # image size related flags
    # assuming that the image has same height and width
    hparams.add_hparam("img_len", 32)
    hparams.add_hparam("num_channels", 3)
    # Local attention params
    hparams.add_hparam("local_and_global_att", False)
    hparams.add_hparam("block_length", 256)
    hparams.add_hparam("block_width", 128)
    hparams.add_hparam("num_encoder_layers", 4)
    hparams.add_hparam("num_decoder_layers", 12)
    hparams.add_hparam("dec_attention_type", cia.AttentionType.LOCAL_1D)
    hparams.add_hparam("block_raster_scan", False)
  
    # multipos attention params
    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++
    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)
    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
  
  
  @registry.register_hparams
  def imagetransformer_base():
    hparams = image_transformer_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_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."""
    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
  
  
  @registry.register_hparams
  def imagetransformer_sep_channels():
    """separate rgb embeddings."""
    hparams = imagetransformer_base()
    hparams.num_heads = 4
    hparams.attention_key_channels = hparams.attention_value_channels = 0
    hparams.hidden_size = 256
    hparams.filter_size = 512
    hparams.num_hidden_layers = 6
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_sep_channels_8l():
    """separate rgb embeddings."""
    hparams = imagetransformer_base()
    hparams.num_heads = 4
    hparams.attention_key_channels = hparams.attention_value_channels = 0
    hparams.hidden_size = 256
    hparams.filter_size = 256
    hparams.num_hidden_layers = 8
    hparams.sampling_method = "random"
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_sep_channels_8l_multipos3():
    """separate rgb embeddings."""
    hparams = imagetransformer_sep_channels_8l()
    hparams.q_filter_width = 3
    hparams.kv_filter_width = 3
    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."""
    hparams = imagetransformer_sep_channels_8l()
    hparams.block_width = 256
    hparams.block_length = 256
    hparams.hidden_size = 512
    hparams.num_heads = 8
    hparams.filter_size = 2048
    hparams.batch_size = 4
    hparams.max_length = 3075
    hparams.layer_preprocess_sequence = "none"
    hparams.layer_postprocess_sequence = "dan"
    hparams.num_decoder_layers = 8
    hparams.layer_prepostprocess_dropout = 0.3
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_base_10l_8h_big_uncond_dr03_dan_64():
    """big 1d model for unconditional generation on imagenet."""
    hparams = imagetransformer_base_10l_8h_big_cond_dr03_dan()
    hparams.unconditional = True
    hparams.max_length = 14000
    hparams.batch_size = 1
    hparams.img_len = 64
    hparams.layer_prepostprocess_dropout = 0.1
    return hparams
  
  
  @registry.register_hparams
  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
    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.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
    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_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
    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_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."""
    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_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."""
    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()
    hparams.block_width = 128
    hparams.block_length = 128
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_base_10l_8h_big_cond_dr03_dan():
    """Best conditional Cifar10 gen param."""
    hparams = imagetransformer_base_8l_8h_big_cond_dr03_dan()
    hparams.num_decoder_layers = 10
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_base_10l_8h_big_uncond_dr03_dan():
    """Best unconditional Cifar10 gen param."""
    hparams = imagetransformer_base_10l_8h_big_cond_dr03_dan()
    hparams.num_decoder_layers = 10
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_base_8l_8h_big_cond_dr03_dan_dilated():
    """Dilated hparams."""
    hparams = imagetransformer_base_8l_8h_big_cond_dr03_dan()
    hparams.gap_sizes = [0, 16, 64, 0, 16, 64, 128, 0]
    hparams.dec_attention_type = cia.AttentionType.DILATED
    hparams.block_length = 128
    hparams.block_width = 128
    hparams.add_hparam("num_memory_blocks", 1)
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_base_8l_8h_big_cond_dr03_dan_dilated_b():
    """Dilated hparams."""
    hparams = imagetransformer_base_8l_8h_big_cond_dr03_dan_dilated()
    hparams.block_width = 64
    hparams.num_memory_blocks = 2
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_base_8l_8h_big_cond_dr03_dan_dilated_c():
    """Dilated hparams."""
    hparams = imagetransformer_base_8l_8h_big_cond_dr03_dan_dilated()
    hparams.block_width = 32
    hparams.num_memory_blocks = 4
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_base_8l_8h_big_cond_dr03_dan_dilated_d():
    """Dilated hparams."""
    hparams = imagetransformer_base_8l_8h_big_cond_dr03_dan_dilated()
    hparams.gap_sizes = [0, 16, 64, 16, 64, 128, 256, 0]
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_base_12l_8h_big():
    """big 1d model for conditional image generation."""
    hparams = imagetransformer_sep_channels_8l_8h()
    hparams.filter_size = 1024
    hparams.num_decoder_layers = 12
    hparams.batch_size = 1
    hparams.hidden_size = 512
    hparams.learning_rate_warmup_steps = 4000
    hparams.sampling_method = "random"
    hparams.beam_size = 1
    hparams.block_width = 256
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer1d_base_8l_64by64():
    """hparams fo 12 layer big 1d model for imagenet 64x64."""
    hparams = image_transformer_base()
    hparams.num_heads = 8
    hparams.hidden_size = 512
    hparams.filter_size = 2048
    hparams.num_decoder_layers = 8
    hparams.batch_size = 1
    hparams.block_length = 512
    hparams.block_width = 768
    hparams.layer_prepostprocess_dropout = 0.1
    hparams.max_length = 14000
    hparams.unconditional = int(False)
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer1d_base_12l_64by64():
    """hparams fo 12 layer big 1d model for imagenet 64x64."""
    hparams = image_transformer_base()
    hparams.num_heads = 8
    hparams.hidden_size = 512
    hparams.filter_size = 2048
    hparams.num_decoder_layers = 12
    hparams.batch_size = 1
    hparams.block_length = 512
    hparams.block_width = 768
    hparams.layer_prepostprocess_dropout = 0.1
    hparams.max_length = 14000
    hparams.unconditional = int(False)
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_base_14l_8h_big():
    """big 1d model for conditional image generation."""
    hparams = imagetransformer_base_12l_8h_big()
    hparams.num_decoder_layers = 14
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_base_14l_8h_big_dr01():
    """big 1d model for conditional image generation."""
    hparams = imagetransformer_base_14l_8h_big()
    hparams.layer_prepostprocess_dropout = 0.1
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_base_12l_8h_big_uncond():
    """big 1d model for conditional image generation."""
    hparams = imagetransformer_base_12l_8h_big()
    hparams.unconditional = True
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_base_14l_8h_big_uncond():
    """big 1d model for conditional image generation."""
    hparams = imagetransformer_base_12l_8h_big_uncond()
    hparams.num_decoder_layers = 14
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_sep_channels_12l_16h_imagenet_large():
    """separate rgb embeddings."""
    hparams = imagetransformer_sep_channels_8l_8h()
    hparams.num_hidden_layers = 12
    hparams.batch_size = 1
    hparams.filter_size = 2048
    hparams.num_heads = 16
    hparams.learning_rate_warmup_steps = 16000
    hparams.sampling_method = "random"
    hparams.learning_rate = 0.1
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_sep_channels_16l_16h_imgnet_lrg_loc():
    """separate rgb embeddings."""
    hparams = imagetransformer_sep_channels_12l_16h_imagenet_large()
    hparams.num_hidden_layers = 16
    hparams.local_attention = True
    hparams.batch_size = 1
    hparams.block_length = 256
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_sep_channels_16l_16h_imgnet_lrg_loc_128():
    """separate rgb embeddings."""
    hparams = imagetransformer_sep_channels_12l_16h_imagenet_large()
    hparams.num_hidden_layers = 16
    hparams.local_attention = True
    hparams.batch_size = 1
    hparams.block_length = 128
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_sep_output_channels_8l_local_and_global_att():
    """separate rgb embeddings."""
    hparams = imagetransformer_sep_channels_8l()
    hparams.sampling_method = "random"
    hparams.local_and_global_att = True
    return hparams
  
  
  @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_dr01()
    # num_hidden_layers
    hparams.num_decoder_layers = 10
    hparams.num_heads = 16
    hparams.hidden_size = 1024
    hparams.filter_size = 4096
    hparams.batch_size = 1
    hparams.layer_prepostprocess_dropout = 0.1
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_base_10l_16h_big_dr01_imgnet():
    """big 1d model for conditional image generation."""
    hparams = imagetransformer_base_14l_8h_big_dr01()
    # num_hidden_layers
    hparams.num_decoder_layers = 10
    hparams.num_heads = 16
    hparams.hidden_size = 1024
    hparams.filter_size = 4096
    hparams.batch_size = 1
    hparams.unconditional = False
    hparams.layer_prepostprocess_dropout = 0.1
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_sep_channels_8l_8h():
    """separate rgb embeddings."""
    hparams = imagetransformer_base()
    hparams.num_heads = 8
    hparams.batch_size = 1
    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"
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_sep_channels_8l_8h_local_and_global_att():
    """separate rgb embeddings."""
    hparams = imagetransformer_sep_channels_8l_8h()
    hparams.num_heads = 8
    hparams.batch_size = 1
    hparams.attention_key_channels = hparams.attention_value_channels = 0
    hparams.hidden_size = 256
    hparams.filter_size = 256
    hparams.num_hidden_layers = 4
    hparams.sampling_method = "random"
    hparams.local_and_global_att = True
    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_dr01()
    # num_hidden_layers
    hparams.num_decoder_layers = 8
    hparams.num_heads = 8
    hparams.hidden_size = 512
    hparams.filter_size = 2048
    hparams.layer_prepostprocess_dropout = 0.3
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_tiny():
    hparams = imagetransformer_base()
    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_tiny()
    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()
    update_hparams_for_tpu(hparams)
    hparams.num_hidden_layers = 2
    hparams.hidden_size = 16
    hparams.batch_size = 2
    hparams.num_heads = 2
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_base_10l_16h_big_dr01_moe_imgnet():
    """big 1d model for conditional image generation."""
    hparams = imagetransformer_base_10l_16h_big_dr01_imgnet()
    hparams.initializer = "orthogonal"
    hparams.learning_rate_warmup_steps = 16000
    hparams.add_hparam("moe_layers_decoder", "2,7")  # Which layer is MoE.
    hparams.moe_hidden_sizes = "4096"  # Hidden layer sizes (comma-separated).
    hparams.moe_num_experts = 64  # Number of experts in each MoE layer.
    hparams.moe_k = 4  # How many experts to use per batch element (try 2 or 4).
    hparams.moe_loss_coef = 3e-2  # MoE loss coefficient (1e-2 is usually ok).
    hparams.scheduled_sampling_prob = 0.1
    hparams.scheduled_sampling_warmup_steps = 200000
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_moe_tiny():
    """Set of hyperparameters for a very small imagetransformer with MoE."""
    hparams = imagetransformer_tiny()
    hparams.hidden_size = 64
    hparams.batch_size = 1
    hparams.num_hidden_layers = 3
    hparams.dec_attention_type = cia.AttentionType.MOE_LOCAL_1D
    hparams.add_hparam("moe_layers_decoder", "1")  # Which layer is MoE.
    hparams.moe_hidden_sizes = "1024"  # Hidden layer sizes (comma-separated).
    hparams.moe_num_experts = 16  # Number of experts in each MoE layer.
    hparams.moe_k = 2  # How many experts to use per batch element (try 2 or 4).
    hparams.moe_loss_coef = 1e-2  # MoE loss coefficient (1e-2 is usually ok).
    return hparams
  
  
  def update_hparams_for_tpu(hparams):
    hparams.optimizer = "Adafactor"
    hparams.learning_rate_schedule = "rsqrt_decay"
    hparams.learning_rate_warmup_steps = 6000
    hparams.batch_size = 4
  
  
  @registry.register_hparams
  def imagetransformer_sep_channels_8l_tpu():
    """Hparams for training imagetransformer on tpu."""
    hparams = imagetransformer_sep_channels_8l()
    update_hparams_for_tpu(hparams)
    hparams.batch_size = 4
    hparams.num_heads = 4   # heads are expensive on tpu
    hparams.shared_embedding_and_softmax_weights = False
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_b10l_4h_big_uncond_dr03_tpu():
    """Small model for tpu 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 = 10
    hparams.block_length = 128
    hparams.hidden_size = 512
    hparams.filter_size = 1024
    hparams.learning_rate = 0.2
    hparams.layer_preprocess_sequence = "none"
    hparams.layer_postprocess_sequence = "dan"
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_b10l_dr03_moe_tpu():
    """Moe tpu params."""
    hparams = imagetransformer_b10l_4h_big_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 = 10
    hparams.layer_preprocess_sequence = "none"
    hparams.layer_postprocess_sequence = "dan"
    hparams.ffn_layer = "local_moe_tpu"
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_b10l_4h_big_uncond_dr03_lr025_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 = 10
    hparams.learning_rate = 0.25
    hparams.learning_rate_warmup_steps = 8000
    hparams.layer_preprocess_sequence = "none"
    hparams.layer_postprocess_sequence = "dan"
    # hparams.unconditional = True
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_b12l_4h_big_uncond_dr03_tpu():
    """TPU 12 layer 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 = 12
    hparams.block_length = 128
    hparams.hidden_size = 512
    hparams.filter_size = 1024
    hparams.layer_preprocess_sequence = "none"
    hparams.layer_postprocess_sequence = "dan"
    hparams.layer_prepostprocess_dropout = 0.3
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_b12l_4h_big_uncond_dr03_lr025_tpu():
    hparams = imagetransformer_b12l_4h_big_uncond_dr03_tpu()
    update_hparams_for_tpu(hparams)
    hparams.learning_rate = 0.25
    hparams.learning_rate_warmup_steps = 5000
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_b12l_4h_b256_uncond_dr03_tpu():
    """works very well on 4x4."""
    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 = 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_b12l_4h_b256_uncond_dr03_rel_tpu():
    """works very well on 4x4."""
    hparams = imagetransformer_b12l_4h_b256_uncond_dr03_tpu()
    hparams.shared_rel = True
    hparams.dec_attention_type = cia.AttentionType.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.AttentionType.RELATIVE_LOCAL_1D, cia.AttentionType.LOCAL_1D])
  
  
  @registry.register_hparams
  def imagetransformer_b12l_4h_b128_h512_uncond_dr03_tpu():
    """TPU related big 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 = 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_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
  
  
  @registry.register_hparams
  def imagetransformer_b12l_4h_uncond_dr03_tpu():
    """TPU related small model."""
    hparams = imagetransformer_b12l_4h_b256_uncond_dr03_tpu()
    hparams.learning_rate = 0.2
    hparams.learning_rate_warmup_steps = 4000
    hparams.layer_preprocess_sequence = "none"
    hparams.layer_postprocess_sequence = "dan"
    hparams.layer_prepostprocess_dropout = 0.3
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_b12l_4h_b128_uncond_dr03_tpu():
    """TPU config for cifar 10."""
    hparams = imagetransformer_bas8l_8h_big_uncond_dr03_imgnet()
    update_hparams_for_tpu(hparams)
    hparams.batch_size = 2
    hparams.num_heads = 4   # heads are expensive on tpu
    hparams.num_decoder_layers = 12
    hparams.block_length = 128
    hparams.hidden_size = 256
    hparams.filter_size = 2048
    hparams.layer_preprocess_sequence = "none"
    hparams.layer_postprocess_sequence = "dan"
    hparams.layer_prepostprocess_dropout = 0.1
    hparams.optimizer = "Adafactor"
    hparams.learning_rate_schedule = "rsqrt_decay"
    hparams.learning_rate_warmup_steps = 10000
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_b12l_8h_b256_uncond_dr03_tpu():
    """TPU related 12 layer 8 heads model."""
    hparams = imagetransformer_bas8l_8h_big_uncond_dr03_imgnet()
    update_hparams_for_tpu(hparams)
    hparams.batch_size = 2
    hparams.num_heads = 8   # heads are expensive on tpu
    hparams.num_decoder_layers = 12
    hparams.block_length = 256
    hparams.hidden_size = 512
    hparams.filter_size = 2048
    hparams.layer_preprocess_sequence = "none"
    hparams.layer_postprocess_sequence = "dan"
    hparams.layer_prepostprocess_dropout = 0.3
    return hparams
  
  
  @registry.register_hparams
  def imagetransformer_b10l_4h_big_uncond_dr01_tpu():
    """big 1d model for conditional image generation."""
    hparams = imagetransformer_b12l_4h_big_uncond_dr03_tpu()
    # num_hidden_layers
    hparams.num_decoder_layers = 10
    hparams.num_heads = 4
    hparams.hidden_size = 1024
    hparams.filter_size = 4096
    hparams.batch_size = 1
    hparams.layer_prepostprocess_dropout = 0.1
    return hparams
  