Range of hyperparameters for vizier.
421899
def imagetransformer_cifar_tpu_range(rhp):
  """"""
  # 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])
634109
def _fused_batch_norm_op(self, input_batch, mean, variance, use_batch_stats):
    """"""
    # Store the original shape of the mean and variance.
    mean_shape = mean.get_shape()
    variance_shape = variance.get_shape()
    # The fused batch norm expects the mean, variance, gamma and beta
    # tensors to have dimension 1, so we flatten them to remove the
    # extra dimensions. In addition, it expects the input_batch to have
    # dimension 4, so we reshape it accordingly.
    gamma_flatten = tf.reshape(self._gamma, shape=(self._num_channels,))
    beta_flatten = tf.reshape(self._beta, shape=(self._num_channels,))
    flatten_mean = tf.reshape(mean, shape=(self._num_channels,))
    flatten_variance = tf.reshape(variance, shape=(self._num_channels,))
    use_batch_stats = tf.convert_to_tensor(use_batch_stats)

    input_shape = input_batch.get_shape()
    output_shape = [-1] + input_shape.as_list()[1:]

    flat_image_size = np.prod(self._image_shape, dtype=np.int32)

    if len(self._data_format) == 4:
      fusable_data_format = self._data_format
      fusable_batch = input_batch
    elif self._channel_index == 1 and self._image_shape:
      fusable_data_format = "NCHW"
      fusable_batch = tf.reshape(
          input_batch,
          shape=(-1, self._num_channels, 1, flat_image_size))
    else:
      # The CPU implementation of FusedBatchNorm only supports NHWC tensor
      # format for now.
      fusable_data_format = "NHWC"
      fusable_batch = tf.reshape(
          input_batch,
          shape=(-1, 1, flat_image_size, self._num_channels))

    common_args = {
        "scale": gamma_flatten,
        "offset": beta_flatten,
        "epsilon": self._eps,
        "data_format": fusable_data_format,
        "name": "batch_norm"
    }

    def use_batch_stats_fused_batch_norm():
      return tf.nn.fused_batch_norm(
          fusable_batch,
          mean=None,
          variance=None,
          is_training=True,
          **common_args)

    def moving_average_fused_batch_norm():
      return tf.nn.fused_batch_norm(
          fusable_batch,
          mean=flatten_mean,
          variance=flatten_variance,
          is_training=False,
          **common_args)

    batch_norm_op, mean, variance = utils.smart_cond(
        use_batch_stats, use_batch_stats_fused_batch_norm,
        moving_average_fused_batch_norm)

    if len(self._data_format) != 4:
      batch_norm_op = tf.reshape(batch_norm_op, output_shape)
    mean = tf.reshape(mean, mean_shape)
    variance = tf.reshape(variance, variance_shape)
    return batch_norm_op, mean, variance
314745
def fit(self,Params0=None,grad_threshold=1e-2):
        

        # GPVD initialization
        lik = limix.CLikNormalNULL()
        # Initial Params
        if Params0==None:
            n_params = self.C1.getNumberParams()
            n_params+= self.C2.getNumberParams()
            Params0 = SP.rand(n_params)
        # MultiGP Framework
        covar = []
        gp    = []
        mean  = []
        for i in range(self.N):
            covar.append(limix.CLinCombCF())
            covar[i].addCovariance(self.C1)
            covar[i].addCovariance(self.C2)
            coeff = SP.array([self.eigen[i],1])
            covar[i].setCoeff(coeff)
            mean.append(limix.CLinearMean(self.Yt[:,i],SP.eye(self.P)))
        gpVD = limix.CGPvarDecomp(covar[0],lik,mean[0],SP.ones(self.N),self.P,self.Yt,Params0)
        for i in range(self.N):
            gp.append(limix.CGPbase(covar[i],lik,mean[i]))
            gp[i].setY(self.Yt[:,i])
            gpVD.addGP(gp[i])
                
        # Optimization
        gpVD.initGPs()
        gpopt = limix.CGPopt(gpVD)
        LML0=-1.0*gpVD.LML()
        start_time = time.time()
        conv = gpopt.opt()
        time_train = time.time() - start_time
        LML=-1.0*gpVD.LML()
        LMLgrad = SP.linalg.norm(gpVD.LMLgrad()['covar'])
        Params = gpVD.getParams()['covar']
    
        # Check whether limix::CVarianceDecomposition.train() has converged
        if conv!=True or LMLgrad>grad_threshold or Params.max()>10:
            print('limix::CVarianceDecomposition::train has not converged')
            res=None
        else:
            res = {
                'Params0':          Params0,
                'Params':           Params,
                'LML':              SP.array([LML]),
                'LML0':             SP.array([LML0]),
                'LMLgrad':          SP.array([LMLgrad]),
                'time_train':       SP.array([time_train]),
                }
        return res
        pass
465040
def BatchNorm(x, params, axis=(0, 1, 2), epsilon=1e-5,
              center=True, scale=True, **unused_kwargs):
  """"""
  mean = np.mean(x, axis, keepdims=True)
  # Fast but less numerically-stable variance calculation than np.var.
  m1 = np.mean(x**2, axis, keepdims=True)
  var = m1 - mean**2
  z = (x - mean) / np.sqrt(var + epsilon)

  # Expand the parameters to have the right axes.
  beta, gamma = params
  # TODO(phawkins): np.expand_dims should accept an axis tuple.
  # (https://github.com/numpy/numpy/issues/12290)
  ed = tuple(None if i in axis else slice(None) for i in range(np.ndim(x)))
  beta = beta[ed]
  gamma = gamma[ed]

  # Return the z rescaled by the parameters if requested.
  if center and scale:
    return gamma * z + beta
  if center:
    return z + beta
  if scale:
    return gamma * z
  return z
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
883989
function boxBlurH(scl, tcl, w, h, r) {
    var iarr = 1 / (r+r+1); // averaging adjustment parameter
    for (var i = 0; i < h; i++) {
        var ti = i * w, // current target index
            li = ti, // current left side of the examined range
            ri = ti + r, // current right side of the examined range
            fv = scl[ti], // first value in the row
            lv = scl[ti + w - 1], // last value in the row
            val = (r+1) * fv, // target value, accumulated over examined points
            j;
        // Sum the source values in the box
        for (j = 0; j < r; j++) { val += scl[ti + j]; }
        // Compute the target value by taking the average of the surrounding
        // values. This is done by adding the deviations so far and adjusting,
        // accounting for the edges by extending the first and last values.
        for (j = 0  ; j <= r ; j++) { val += scl[ri++] - fv       ; tcl[ti++] = val*iarr; }
        for (j = r+1; j < w-r; j++) { val += scl[ri++] - scl[li++]; tcl[ti++] = val*iarr; }
        for (j = w-r; j < w  ; j++) { val += lv        - scl[li++]; tcl[ti++] = val*iarr; }
    }
}
946108
def bytenet_base():
  """"""
  hparams = common_hparams.basic_params1()
  hparams.batch_size = 2048
  hparams.hidden_size = 768
  hparams.dropout = 0.2
  hparams.symbol_dropout = 0.2
  hparams.label_smoothing = 0.1
  hparams.clip_grad_norm = 2.0
  hparams.num_hidden_layers = 4
  hparams.kernel_height = 3
  hparams.kernel_width = 1
  hparams.learning_rate_decay_scheme = "exp"
  hparams.learning_rate = 0.05
  hparams.learning_rate_warmup_steps = 3000
  hparams.initializer_gain = 1.0
  hparams.weight_decay = 3.0
  hparams.num_sampled_classes = 0
  hparams.sampling_method = "argmax"
  hparams.optimizer_adam_epsilon = 1e-6
  hparams.optimizer_adam_beta1 = 0.85
  hparams.optimizer_adam_beta2 = 0.997
  hparams.add_hparam("num_block_repeat", 4)
  return hparams
1215809
def vqa_attention_base():
  """"""
  hparams = common_hparams.basic_params1()
  hparams.batch_size = 128
  hparams.use_fixed_batch_size = True,
  hparams.optimizer = "adam"
  hparams.optimizer_adam_beta1 = 0.9
  hparams.optimizer_adam_beta2 = 0.999
  hparams.optimizer_adam_epsilon = 1e-8
  hparams.weight_decay = 0.
  hparams.clip_grad_norm = 0.
  hparams.initializer = "xavier"
  hparams.learning_rate = 0.5
  hparams.learning_rate_schedule = "legacy"
  hparams.learning_rate_warmup_steps = 0
  hparams.learning_rate_decay_scheme = "exp"
  hparams.learning_rate_decay_rate = 0.5
  hparams.learning_rate_decay_steps = 50000
  hparams.dropout = 0.5
  hparams.summarize_grads = True
  hparams.summarize_vars = True

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

  # add new hparams
  # preprocess
  hparams.add_hparam("resize_side", 512)
  hparams.add_hparam("height", 448)
  hparams.add_hparam("width", 448)
  hparams.add_hparam("distort", True)

  hparams.add_hparam("train_resnet", False)
  hparams.add_hparam("rnn_type", "lstm")
  hparams.add_hparam("num_rnn_layers", 1)
  hparams.add_hparam("max_question_length", 15)
  # lstm hidden size
  hparams.hidden_size = 512

  hparams.add_hparam("attn_dim", 512)
  hparams.add_hparam("num_glimps", 2)

  hparams.add_hparam("num_mlp_layers", 1)
  hparams.add_hparam("mlp_dim", 1024)

  hparams.add_hparam("image_input_type", "image")
  hparams.add_hparam("image_model_fn", "resnet_v1_152")
  hparams.add_hparam("image_feat_size", 0)

  # self attention parts
  hparams.norm_type = "layer"
  hparams.layer_preprocess_sequence = "n"
  hparams.layer_postprocess_sequence = "da"
  hparams.layer_prepostprocess_dropout = 0.3
  hparams.attention_dropout = 0.1
  hparams.relu_dropout = 0.1
  hparams.image_hidden_size = 2048
  hparams.add_hparam("num_encoder_layers", 1)
  # Attention-related flags.
  hparams.add_hparam("num_heads", 8)
  hparams.add_hparam("attention_key_channels", 0)
  hparams.add_hparam("attention_value_channels", 0)
  hparams.add_hparam("image_filter_size", 1024)
  hparams.add_hparam("self_attention_type", "dot_product")
  hparams.add_hparam("scale_dotproduct", True)

  return hparams
923233
def hyperparameters(self):
        """"""
        hyperparameters = super(TensorFlow, self).hyperparameters()

        self.checkpoint_path = self.checkpoint_path or self._default_s3_path('checkpoints')
        mpi_enabled = False

        if self._script_mode_enabled():
            additional_hyperparameters = {}

            if 'parameter_server' in self.distributions:
                ps_enabled = self.distributions['parameter_server'].get('enabled', False)
                additional_hyperparameters[self.LAUNCH_PS_ENV_NAME] = ps_enabled

            if 'mpi' in self.distributions:
                mpi_dict = self.distributions['mpi']
                mpi_enabled = mpi_dict.get('enabled', False)
                additional_hyperparameters[self.LAUNCH_MPI_ENV_NAME] = mpi_enabled
                additional_hyperparameters[self.MPI_NUM_PROCESSES_PER_HOST] = mpi_dict.get('processes_per_host', 1)
                additional_hyperparameters[self.MPI_CUSTOM_MPI_OPTIONS] = mpi_dict.get('custom_mpi_options', '')

            self.model_dir = self.model_dir or self._default_s3_path('model', mpi=mpi_enabled)
            additional_hyperparameters['model_dir'] = self.model_dir
        else:
            additional_hyperparameters = {'checkpoint_path': self.checkpoint_path,
                                          'training_steps': self.training_steps,
                                          'evaluation_steps': self.evaluation_steps,
                                          'sagemaker_requirements': self.requirements_file}

        hyperparameters.update(Framework._json_encode_hyperparameters(additional_hyperparameters))
        return hyperparameters
1087411
def transformer_base_single_gpu():
  """"""
  hparams = transformer_base()
  hparams.batch_size = 1024
  hparams.learning_rate_schedule = "constant*linear_warmup*rsqrt_decay"
  hparams.learning_rate_constant = 0.1
  hparams.learning_rate_warmup_steps = 16000
  return hparams
568896
def vizarray(x, cmap=None, scale=None, vmin=None, vmax=None, block_size=None):
    """"""
    if not (x.ndim == 2 or x.ndim == 1):
        raise TypeError('This function only works with 1 or 2 dimensional arrays')
    global _cmap, _scale, _vmin, _vmax, _block_size
    cmap = cmap if cmap is not None else _cmap
    scale = scale if scale is not None else _scale
    vmin = vmin if vmin is not None else _vmin
    vmax = vmax if vmax is not None else _vmax
    block_size = block_size if block_size is not None else _block_size
    base = x.base if x.base is not None else None
    data = x.copy()
    if scale:
        n = colors.Normalize(vmin=vmin, vmax=vmax)
        if base is not None:
            n.autoscale(base)
        data = n(data)
    if data.ndim == 1:
        rows = 1
        cols = data.shape[0]
        bg = BlockGrid(cols, rows, block_size=block_size)
        for col in range(cols):
            bg[0,col] = _value_to_color(data[col], cmap)
    elif data.ndim == 2:
        rows = data.shape[0]
        cols = data.shape[1]
        bg = BlockGrid(cols, rows, block_size=block_size)
        for row in range(rows):
            for col in range(cols):
                bg[row, col] = _value_to_color(data[row, col], cmap)
    return bg
1082129
def image_encoder(image_feat,
                  hparams,
                  name="image_encoder",
                  save_weights_to=None,
                  make_image_summary=True):
  """"""

  x = image_feat
  image_hidden_size = hparams.image_hidden_size or hparams.hidden_size
  image_filter_size = hparams.image_filter_size or hparams.filter_size
  with tf.variable_scope(name):
    for layer in range(hparams.num_encoder_layers or hparams.num_hidden_layers):
      with tf.variable_scope("layer_%d" % layer):
        with tf.variable_scope("self_attention"):
          y = vqa_layers.multihead_attention(
              common_layers.layer_preprocess(x, hparams),
              None,
              None,
              hparams.attention_key_channels or image_hidden_size,
              hparams.attention_value_channels or image_hidden_size,
              image_hidden_size,
              hparams.num_heads,
              hparams.attention_dropout,
              attention_type=hparams.image_self_attention_type,
              save_weights_to=save_weights_to,
              make_image_summary=make_image_summary,
              scale_dotproduct=hparams.scale_dotproduct,
          )
          utils.collect_named_outputs(
              "norms", "image_feat_self_attention_%d"%(layer),
              tf.norm(y, axis=-1))
          x = common_layers.layer_postprocess(x, y, hparams)
          utils.collect_named_outputs(
              "norms", "image_feat_self_attention_postprocess_%d"%(layer),
              tf.norm(x, axis=-1))
        with tf.variable_scope("ffn"):
          y = common_layers.dense_relu_dense(
              common_layers.layer_preprocess(x, hparams),
              image_filter_size,
              image_hidden_size,
              dropout=hparams.relu_dropout,
          )
          utils.collect_named_outputs(
              "norms", "image_feat_ffn_%d"%(layer), tf.norm(y, axis=-1))
          x = common_layers.layer_postprocess(x, y, hparams)
          utils.collect_named_outputs(
              "norms", "image_feat_ffn_postprocess_%d"%(layer),
              tf.norm(x, axis=-1))
    # if normalization is done in layer_preprocess, then it should also be done
    # on the output, since the output can grow very large, being the sum of
    # a whole stack of unnormalized layer outputs.
    return common_layers.layer_preprocess(x, hparams)
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
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
1111755
def mtf_image_transformer_base_cifar():
  """"""
  hparams = mtf_image_transformer_base()
  hparams.mesh_shape = "batch:8"
  hparams.layout = "batch:batch"
  hparams.learning_rate_decay_steps = 13600  # one epoch
  hparams.batch_size = 32
  hparams.num_heads = 4
  hparams.num_decoder_layers = 12
  hparams.block_length = 256
  hparams.hidden_size = 512
  hparams.d_ff = 2048
  hparams.learning_rate = 0.5
  hparams.layer_preprocess_sequence = "none"
  hparams.layer_postprocess_sequence = "dan"
  hparams.layer_prepostprocess_dropout = 0.3
  hparams.unconditional = True
  return hparams
619619
def mtf_image_transformer_base_imagenet_mp128():
  """"""
  hparams = mtf_image_transformer_base_imagenet()
  hparams.mesh_shape = "model:8;batch:4"
  hparams.layout = "batch:batch;d_ff:model;heads:model"
  hparams.batch_size = 8
  hparams.img_len = 128
  hparams.block_length = 128
  hparams.num_heads = 8
  hparams.num_decoder_layers = 4
  hparams.d_ff = 4096
  hparams.learning_rate_warmup_steps = 31250
  hparams.unconditional = True
  hparams.max_length = 256*256*3
  return hparams
640
def slicenet_range1(ranged_hparams):
  """"""
  rhp = ranged_hparams
  rhp.set_float("clip_grad_norm", 1.0, 10.0, scale=rhp.LOG_SCALE)
  rhp.set_float("learning_rate", 0.02, 1.0, scale=rhp.LOG_SCALE)
  rhp.set_float("optimizer_adam_beta2", 0.995, 0.998)
  rhp.set_float("weight_decay", 1.0, 5.0)