HParams for training image_cifar10_plain_gen_flat_rev with memory.
355335
def transformer_cifar10_memory_v0():
  """"""
  hparams = transformer_wikitext103_l4k_memory_v0()

  hparams.num_hidden_layers = 6

  hparams.max_length = 32 * 32 * 3
  hparams.split_targets_chunk_length = 64 * 3
  hparams.split_targets_max_chunks = int(
      hparams.max_length / hparams.split_targets_chunk_length)
  hparams.num_memory_items = 128 * 3

  # Since this is an image problem, batch size refers to examples (not tokens)
  target_images_per_batch = 4
  hparams.batch_size = int(target_images_per_batch * (
      hparams.max_length / hparams.split_targets_chunk_length))

  # The recurrent memory needs to know the actual batch size (in sequences)
  hparams.recurrent_memory_batch_size = hparams.batch_size

  hparams.max_relative_position = (
      hparams.num_memory_items + hparams.split_targets_chunk_length)

  return hparams
810724
def mtf_bitransformer_base():
  """"""
  hparams = mtf_transformer2_base()
  hparams.max_length = 256
  hparams.shared_embedding = True
  # HYPERPARAMETERS FOR THE LAYER STACKS
  hparams.add_hparam("encoder_layers", ["self_att", "drd"] * 6)
  hparams.add_hparam("decoder_layers", ["self_att", "enc_att", "drd"] * 6)
  hparams.add_hparam("encoder_num_layers", 6)
  hparams.add_hparam("decoder_num_layers", 6)
  # number of heads in multihead attention
  hparams.add_hparam("encoder_num_heads", 8)
  hparams.add_hparam("decoder_num_heads", 8)
  hparams.add_hparam("local_attention_radius", 128)

  # default of 0 for standard transformer behavior
  # 1 means a single set of keys and values that are read by all query heads
  hparams.add_hparam("encoder_num_memory_heads", 0)
  hparams.add_hparam("decoder_num_memory_heads", 0)
  # share attention keys and values
  hparams.add_hparam("encoder_shared_kv", False)
  hparams.add_hparam("decoder_shared_kv", False)

  # Parameters for computing the maximum decode length in beam search.
  # Maximum decode length is:
  #    min(max_length,
  #        decode_length_multiplier * input_length + decode_length_constant)
  hparams.add_hparam("decode_length_multiplier", 1.5)
  hparams.add_hparam("decode_length_constant", 10.0)
  # used during decoding
  hparams.add_hparam("alpha", 0.6)
  hparams.sampling_temp = 0.0
  return hparams
7927
def add_standard_attention_hparams(hparams):
  """"""
  # All hyperparameters ending in "dropout" are automatically set to 0.0
  # when not in training mode.

  # hparams used and which should have been defined outside (in
  # common_hparams):
  # Global flags
  # hparams.mode
  # hparams.hidden_size
  # Pre-post processing flags
  # hparams.layer_preprocess_sequence
  # hparams.layer_postprocess_sequence
  # hparams.layer_prepostprocess_dropout
  # hparams.norm_type
  # hparams.norm_epsilon
  # Mixture-of-Expert flags
  # hparams.moe_hidden_sizes
  # hparams.moe_num_experts
  # hparams.moe_k
  # hparams.moe_loss_coef

  # Attention layers flags
  hparams.add_hparam("num_heads", 8)
  hparams.add_hparam("attention_key_channels", 0)
  hparams.add_hparam("attention_value_channels", 0)
  hparams.add_hparam("attention_dropout", 0.0)
  # Attention: Local
  hparams.add_hparam("attention_loc_block_length", 256)
  # Attention: Local (unmasked only): How much to look left.
  hparams.add_hparam("attention_loc_block_width", 128)
  # Attention: Memory-compressed
  hparams.add_hparam("attention_red_factor", 3)
  hparams.add_hparam("attention_red_type", "conv")
  hparams.add_hparam("attention_red_nonlinearity", "none")

  # Fully connected layers flags
  # To be more consistent, should use filter_size to also control the MOE
  # size if moe_hidden_sizes not set.
  hparams.add_hparam("filter_size", 2048)
  hparams.add_hparam("relu_dropout", 0.0)

  return hparams
356926
def imagetransformer_sep_channels_8l_8h():
  """"""
  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
126757
def adafactor_optimizer_from_hparams(hparams, lr):
  
  if hparams.optimizer_adafactor_decay_type == "Adam":
    decay_rate = adafactor_decay_rate_adam(
        hparams.optimizer_adafactor_beta2)
  elif hparams.optimizer_adafactor_decay_type == "pow":
    decay_rate = adafactor_decay_rate_pow(
        hparams.optimizer_adafactor_memory_exponent)
  else:
    raise ValueError("unknown optimizer_adafactor_decay_type")
  return AdafactorOptimizer(
      multiply_by_parameter_scale=(
          hparams.optimizer_adafactor_multiply_by_parameter_scale),
      learning_rate=lr,
      decay_rate=decay_rate,
      beta1=hparams.optimizer_adafactor_beta1,
      clipping_threshold=hparams.optimizer_adafactor_clipping_threshold,
      factored=hparams.optimizer_adafactor_factored)
1339652
def transformer_symshard_base():
  """"""
  hparams = common_hparams.basic_params1()
  hparams.hidden_size = 256
  hparams.batch_size = 2048
  hparams.max_length = 0
  # All hyperparameters ending in "dropout" are automatically set to 0.0
  # when not in training mode.
  hparams.layer_prepostprocess_dropout = 0.2
  hparams.add_hparam("attention_dropout", 0.1)
  hparams.add_hparam("relu_dropout", 0.0)
  hparams.add_hparam("relu_dropout_broadcast_dims", "1")
  hparams.layer_prepostprocess_dropout = 0.1
  hparams.layer_prepostprocess_dropout_broadcast_dims = "1"  # length
  hparams.label_smoothing = 0.1
  hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
  hparams.optimizer = "Adafactor"
  hparams.learning_rate_schedule = "rsqrt_decay"
  hparams.learning_rate_warmup_steps = 10000
  hparams.initializer_gain = 1.0
  hparams.initializer = "uniform_unit_scaling"
  hparams.weight_decay = 0.0
  # TODO(noam): use this to control sharing.  We now share always
  hparams.shared_embedding_and_softmax_weights = True
  # we only want one data shard.
  hparams.no_data_parallelism = True
  # bypass the symbol modality so that we can use model parallelism.
  hparams.bottom = {
      "inputs": modalities.identity_bottom,
      "targets": modalities.identity_bottom,
  }
  hparams.top = {
      "targets": modalities.identity_top,
  }
  hparams.add_hparam("filter_size", 1280)
  hparams.add_hparam("mix_fraction", 0.5)
  # attention-related flags
  hparams.add_hparam("multihead_attention_num_heads", 4)
  hparams.add_hparam("multihead_attention_key_channels", 0)
  hparams.add_hparam("multihead_attention_value_channels", 0)
  hparams.add_hparam("pos", "timing")  # timing, none
  hparams.add_hparam(
      "encoder_layers", ("n,att,m,d,a," "n,ffn,m,d,a,") * 6 + "n,d")
  hparams.add_hparam(
      "decoder_layers",
      ("n,att,m,d,a," "n,enc-att,m,d,a," "n,ffn,m,d,a,") * 6 + "n,d")
  # Number of model shards - each one has separate parameters.
  # Changing this number invalidates checkpoints.
  hparams.add_hparam("num_model_shards", 8)
  return hparams
301757
def mtf_image_transformer_base_imagenet():
  """"""
  hparams = mtf_image_transformer_base_cifar()
  hparams.mesh_shape = "batch:32"
  hparams.layout = "batch:batch"
  hparams.batch_size = 128
  hparams.d_ff = 2048
  hparams.hidden_size = 512
  hparams.num_decoder_layers = 12
  hparams.learning_rate = 0.5
  hparams.learning_rate_warmup_steps = 31250
  hparams.layer_preprocess_sequence = "none"
  hparams.layer_postprocess_sequence = "dan"
  hparams.layer_prepostprocess_dropout = 0.1
  hparams.unconditional = True
  return hparams
830593
def imagetransformer_b10l_4h_big_uncond_dr03_lr025_tpu():
  """"""
  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
771227
def imagetransformer_ae_cifar():
  """"""
  hparams = transformer_ae_small()
  hparams.filter_size = 512
  hparams.num_compress_steps = 3
  hparams.startup_steps = 10000
  hparams.is_2d = 0
  hparams.learning_rate_warmup_steps = 8000
  hparams.learning_rate = 0.2
  hparams.hidden_size = 512
  hparams.batch_size = 1
  hparams.max_length = 256
  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.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.norm_type = "layer"
  hparams.layer_prepostprocess_dropout = 0.0
  hparams.num_heads = 8
  hparams.task = "image"
  hparams.ffn_layer = "conv_hidden_relu"
  # All hyperparameters ending in "dropout" are automatically set to 0.0
  # when not in training mode.
  hparams.attention_dropout = 0.0
  hparams.relu_dropout = 0.
  hparams.pos = "timing"  # timing, none
  hparams.nbr_decoder_problems = 1
  hparams.num_output_layers = 3
  # TODO(trandustin): semhash doesn't work if filter_size != hidden_size. For
  # now, set default to dvq.
  hparams.bottleneck_kind = "dvq"
  hparams.add_hparam("block_size", 1)

  # dilated attention based flags
  hparams.add_hparam("gap_sizes", [2, 4, 8, 16, 32, 64, 2, 4, 8, 16, 32, 64])
  hparams.add_hparam("dilated_attention", False)

  # 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.num_encoder_layers = 4
  hparams.num_decoder_layers = 12
  hparams.add_hparam("dec_attention_type", cia.AttentionType.LOCAL_1D)
  hparams.add_hparam("block_raster_scan", False)
  hparams.add_hparam("shared_rel", False)

  # multipos attention params
  hparams.add_hparam("q_filter_width", 1)
  hparams.add_hparam("kv_filter_width", 1)

  hparams.add_hparam("unconditional", False)  # unconditional generation

  hparams.bottom["targets"] = modalities.image_channel_embeddings_bottom
  hparams.top["targets"] = modalities.image_channel_embeddings_top
  hparams.drop_inputs = True
  hparams.do_attend_compress = False
  hparams.do_attend_decompress = False
  return hparams
1349023
def mtf_resnet_base():
  """"""
  hparams = common_hparams.basic_params1()
  hparams.no_data_parallelism = True
  hparams.use_fixed_batch_size = True
  hparams.batch_size = 32
  hparams.max_length = 3072
  hparams.hidden_size = 256
  hparams.label_smoothing = 0.0
  # 8-way model-parallelism
  hparams.add_hparam("mesh_shape", "batch:8")
  hparams.add_hparam("layout", "batch:batch")
  hparams.add_hparam("filter_size", 1024)

  hparams.add_hparam("num_layers", 6)
  # Share weights between input and target embeddings
  hparams.shared_embedding = True

  hparams.shared_embedding_and_softmax_weights = True
  hparams.optimizer = "Adafactor"
  hparams.learning_rate_schedule = "rsqrt_decay"
  hparams.learning_rate_warmup_steps = 10000
  hparams.add_hparam("d_kv", 32)

  # Image related hparams
  hparams.add_hparam("img_len", 32)
  hparams.add_hparam("num_channels", 3)
  hparams.add_hparam("row_blocks", 1)
  hparams.add_hparam("col_blocks", 1)
  hparams.add_hparam("rows_size", 32)
  hparams.add_hparam("cols_size", 32)

  # Model-specific parameters
  hparams.add_hparam("layer_sizes", [3, 4, 6, 3])
  hparams.add_hparam("filter_sizes", [64, 64, 128, 256, 512])
  hparams.add_hparam("is_cifar", False)

  # Variable init
  hparams.initializer = "normal_unit_scaling"
  hparams.initializer_gain = 2.

  # TODO(nikip): Change optimization scheme?
  hparams.learning_rate = 0.1
  return hparams
570747
def basic_range1(ranged_hparams):
  """"""
  rhp = ranged_hparams
  rhp.set_discrete("batch_size", [1024, 2048, 4096])
  rhp.set_discrete("num_hidden_layers", [1, 2, 3, 4, 5, 6])
  rhp.set_discrete("hidden_size", [32, 64, 128, 256, 512], scale=rhp.LOG_SCALE)
  rhp.set_discrete("kernel_height", [1, 3, 5, 7])
  rhp.set_discrete("kernel_width", [1, 3, 5, 7])
  rhp.set_discrete("compress_steps", [0, 1, 2])
  rhp.set_float("dropout", 0.0, 0.5)
  rhp.set_float("weight_decay", 1e-4, 10.0, scale=rhp.LOG_SCALE)
  rhp.set_float("label_smoothing", 0.0, 0.2)
  rhp.set_float("clip_grad_norm", 0.01, 50.0, scale=rhp.LOG_SCALE)
  rhp.set_float("learning_rate", 0.005, 2.0, scale=rhp.LOG_SCALE)
  rhp.set_categorical("initializer",
                      ["uniform", "orthogonal", "uniform_unit_scaling"])
  rhp.set_float("initializer_gain", 0.5, 3.5)
  rhp.set_categorical("learning_rate_decay_scheme",
                      ["none", "sqrt", "noam", "exp"])
  rhp.set_float("optimizer_adam_epsilon", 1e-7, 1e-2, scale=rhp.LOG_SCALE)
  rhp.set_float("optimizer_adam_beta1", 0.8, 0.9)
  rhp.set_float("optimizer_adam_beta2", 0.995, 0.999)
  rhp.set_categorical(
      "optimizer",
      ["adam", "adagrad", "momentum", "rms_prop", "sgd", "yellow_fin"])
82338
def vqa_self_attention_base():
  """"""
  hparams = common_hparams.basic_params1()
  hparams.batch_size = 128
  hparams.use_fixed_batch_size = True,
  hparams.optimizer = "adam"
  hparams.optimizer_adam_beta1 = 0.9
  hparams.optimizer_adam_beta2 = 0.997
  hparams.optimizer_adam_epsilon = 1e-9
  hparams.weight_decay = 0.
  hparams.clip_grad_norm = 0.
  hparams.initializer = "xavier"
  hparams.learning_rate_schedule = (
      "constant*linear_warmup*rsqrt_normalized_decay")
  hparams.learning_rate_warmup_steps = 8000
  hparams.learning_rate_constant = 1e-3
  hparams.learning_rate_decay_rate = 0.5
  hparams.learning_rate_decay_steps = 50000
  hparams.dropout = 0.5
  hparams.summarize_grads = True
  hparams.summarize_vars = True

  # not used hparams
  hparams.label_smoothing = 0.
  hparams.multiply_embedding_mode = "sqrt_depth"

  # add new hparams
  # use raw image as input
  hparams.add_hparam("image_input_type", "image")
  hparams.add_hparam("image_model_fn", "resnet_v1_152")
  hparams.add_hparam("resize_side", 512)
  hparams.add_hparam("height", 448)
  hparams.add_hparam("width", 448)
  hparams.add_hparam("distort", True)
  hparams.add_hparam("train_resnet", False)

  # image parts
  hparams.add_hparam("image_feat_preprocess_proj", True)
  hparams.add_hparam("image_feat_preprocess_layernorm", True)
  hparams.add_hparam("image_feat_encode", True)
  hparams.add_hparam("image_hidden_size", 0)  # default to hidden_size
  hparams.add_hparam("image_filter_size", 0)  # defaults to filter_size

  # question hidden size
  hparams.hidden_size = 512
  hparams.filter_size = 1024
  hparams.num_hidden_layers = 4

  hparams.add_hparam("multimodal_combine", "concat")
  hparams.add_hparam("num_mlp_layers", 1)
  hparams.add_hparam("mlp_size", 1024)

  # self attention parts
  hparams.norm_type = "layer"
  hparams.layer_preprocess_sequence = "n"
  hparams.layer_postprocess_sequence = "da"
  hparams.layer_prepostprocess_dropout = 0.1
  hparams.attention_dropout = 0.1
  hparams.relu_dropout = 0.1
  hparams.add_hparam("pos", "timing")
  hparams.add_hparam("num_encoder_layers", 0)
  hparams.add_hparam("num_decoder_layers", 0)
  hparams.add_hparam("num_heads", 8)
  hparams.add_hparam("attention_key_channels", 0)
  hparams.add_hparam("attention_value_channels", 0)
  hparams.add_hparam("self_attention_type", "dot_product")
  hparams.add_hparam("image_self_attention_type", "dot_product")
  hparams.add_hparam("question_self_attention_type", "dot_product")
  hparams.add_hparam("block_length", 1)
  hparams.add_hparam("scale_dotproduct", True)

  # iterative part
  hparams.add_hparam("num_rec_steps", 3)

  return hparams
1338359
def imagetransformer_b12l_4h_b256_uncond_dr03_tpu():
  """"""
  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
856793
def imagetransformer_sep_channels():
  """"""
  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
12161
def aligned_base():
  
  hparams = common_hparams.basic_params1()
  hparams.hidden_size = 512
  hparams.batch_size = 5000
  hparams.max_length = 0
  hparams.min_length_bucket = 1024
  hparams.dropout = 0.0
  hparams.layer_prepostprocess_dropout = 0.0
  hparams.label_smoothing = 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 = 2000
  hparams.initializer_gain = 1.0
  hparams.initializer = "uniform_unit_scaling"
  hparams.weight_decay = 0.0
  hparams.optimizer_adam_beta1 = 0.9
  hparams.optimizer_adam_beta2 = 0.98
  hparams.shared_embedding_and_softmax_weights = True
  hparams.add_hparam("ffn_hidden_sizes", "2048")  # Add new ones like this.
  hparams.moe_num_experts = 32
  hparams.layer_preprocess_sequence = "n"
  hparams.layer_postprocess_sequence = "da"
  hparams.add_hparam("layers", "timing," + "conv,att,ffn," * 2)

  # attention-related flags
  hparams.add_hparam("num_heads", 8)
  hparams.add_hparam("attention_key_channels", 0)
  hparams.add_hparam("attention_value_channels", 0)
  # 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("pos", "timing")  # timing, none
  # moe params. local attention moe.
  hparams.add_hparam("attention_local", False)
  hparams.add_hparam("attention_moe_k", 2)
  hparams.add_hparam("attention_num_experts", 16)
  hparams.add_hparam("attention_split_batch", False)
  # Key, query and value dimensions for the attention
  hparams.add_hparam("attention_kq_size", 128)
  hparams.add_hparam("attention_v_size", 256)
  # Loss coef for load balancing
  hparams.add_hparam("attention_load_balance", 2e-2)
  hparams.add_hparam("diet_experts", False)
  hparams.add_hparam("memory_efficient_ffn", False)
  hparams.add_hparam("local_attention_window", 128)
  hparams.add_hparam("attention_num_groups", 8)
  hparams.add_hparam("memory_target_density", 2.0)
  hparams.add_hparam("multiplicative_overhead", 1.25)
  hparams.add_hparam("multiplicative_overhead_eval", 2.0)
  hparams.add_hparam("attention_image_summary", True)
  # LSH params
  hparams.add_hparam("lsh_truncated", True)
  # For testing right-masking.
  # This is not implemented in all layers.
  hparams.add_hparam("mask_right", False)
  return hparams
693341
def mtf_image_transformer_cifar_mp_4x():
  """"""
  hparams = mtf_image_transformer_base_cifar()
  hparams.mesh_shape = "model:4;batch:8"
  hparams.layout = "batch:batch;d_ff:model;heads:model"
  hparams.batch_size = 32
  hparams.num_heads = 8
  hparams.d_ff = 8192
  return hparams
1300807
def attention_lm_moe_base():
  
  hparams = common_hparams.basic_params1()
  hparams.hidden_size = 1024
  hparams.batch_size = 8192
  hparams.max_length = 256
  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 = 2000
  hparams.initializer_gain = 1.0
  hparams.num_hidden_layers = 4
  hparams.initializer = "uniform_unit_scaling"
  hparams.weight_decay = 0.0
  hparams.optimizer_adam_beta1 = 0.9
  hparams.optimizer_adam_beta2 = 0.98
  hparams.num_sampled_classes = 0
  hparams.label_smoothing = 0.0
  hparams.shared_embedding_and_softmax_weights = False
  hparams.add_hparam("filter_size", 2048)  # Add new ones like this.
  hparams.moe_num_experts = 32
  # attention-related flags
  hparams.add_hparam("num_heads", 8)
  hparams.add_hparam("attention_key_channels", 0)
  hparams.add_hparam("attention_value_channels", 0)
  # 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("moe_layers", "2")  # comma separated list of layer numbers
  # moe params. local attention moe.
  # If attention_layers is set, the num_hidden_layers parameter will be ignored
  # and each caracter of the string will correspond to one attention
  # layer type
  hparams.add_hparam("attention_layers", "")
  hparams.add_hparam("attention_type", AttentionType.MULTIHEAD)
  hparams.add_hparam("attention_local", False)
  hparams.add_hparam("attention_moe_k", 2)
  hparams.add_hparam("attention_num_head", 1)
  hparams.add_hparam("attention_num_experts", 16)
  hparams.add_hparam("attention_split_batch", False)
  hparams.add_hparam("attention_red_factor", 3)
  hparams.add_hparam("attention_block_length", 128)
  hparams.add_hparam("attention_reduction_type", "conv")
  # Non linearity for the attention reduction. Either "none", or "silu" (
  # Sigmoid Linear-Unit described in https://arxiv.org/abs/1710.05941)
  hparams.add_hparam("attention_nonlinearity", "none")
  # If attention_exp_factor is set, each input to local_expert_attention (of
  # dimensionality hidden size) is projected into attention_exp_factor smaller
  # inputs, each of dimensionality attention_exp_inputdim. (otherwise
  # attention_exp_inputdim is ignored)
  hparams.add_hparam("attention_exp_factor", 0)
  hparams.add_hparam("attention_exp_inputdim", 128)
  # Key, query and value dimensions for the attention
  hparams.add_hparam("attention_kq_size", 128)
  hparams.add_hparam("attention_v_size", 256)
  # Loss coef for load balancing
  hparams.add_hparam("attention_load_balance", 2e-2)
  # Locality-sensitive hashing params
  hparams.add_hparam("lsh_num_hyperplanes", 4)
  hparams.add_hparam("lsh_use_map_fn", False)

  hparams.add_hparam("use_sepconv", False)
  hparams.add_hparam("diet_experts", False)
  hparams.add_hparam("memory_efficient_ffn", False)
  # if True, we learn a non-autoregressive model from "inputs" to "targets".
  # if False, we learn an autoregressive model to generate "targets"
  hparams.add_hparam("use_inputs", False)
  return hparams