Hyperparameters for CIFAR-10 experiments.
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
1020602
def vq_discrete_bottleneck(x, hparams):
  """"""
  tf.logging.info("Using EMA with beta = {}".format(hparams.beta))
  bottleneck_size = 2**hparams.bottleneck_bits
  x_shape = common_layers.shape_list(x)
  x = tf.reshape(x, [-1, hparams.hidden_size])
  x_means_hot, e_loss = vq_nearest_neighbor(
      x, hparams)
  means, ema_means, ema_count = (hparams.means, hparams.ema_means,
                                 hparams.ema_count)

  # Update the ema variables
  updated_ema_count = moving_averages.assign_moving_average(
      ema_count,
      tf.reduce_sum(x_means_hot, axis=0),
      hparams.decay,
      zero_debias=False)

  dw = tf.matmul(x_means_hot, x, transpose_a=True)
  updated_ema_means = moving_averages.assign_moving_average(
      ema_means, dw, hparams.decay, zero_debias=False)
  n = tf.reduce_sum(updated_ema_count, axis=-1, keepdims=True)
  updated_ema_count = (
      (updated_ema_count + hparams.epsilon) /
      (n + bottleneck_size * hparams.epsilon) * n)
  # pylint: disable=g-no-augmented-assignment
  updated_ema_means = updated_ema_means / tf.expand_dims(
      updated_ema_count, axis=-1)
  # pylint: enable=g-no-augmented-assignment
  with tf.control_dependencies([e_loss]):
    update_means = tf.assign(means, updated_ema_means)
    with tf.control_dependencies([update_means]):
      loss = hparams.beta * e_loss

  discrete = tf.reshape(x_means_hot, x_shape[:-1] + [bottleneck_size])
  return discrete, loss
1367122
def parametrized_bottleneck(x, hparams):
  """"""
  if hparams.bottleneck_kind == "tanh_discrete":
    d, _ = tanh_discrete_bottleneck(
        x, hparams.bottleneck_bits, hparams.bottleneck_noise * 0.5,
        hparams.discretize_warmup_steps, hparams.mode)
    return d, 0.0
  if hparams.bottleneck_kind == "isemhash":
    return isemhash_bottleneck(
        x, hparams.bottleneck_bits, hparams.bottleneck_noise * 0.5,
        hparams.discretize_warmup_steps, hparams.mode,
        hparams.isemhash_noise_dev, hparams.isemhash_mix_prob)
  if hparams.bottleneck_kind == "vq":
    return vq_discrete_bottleneck(x, hparams.bottleneck_bits, hparams.vq_beta,
                                  hparams.vq_decay, hparams.vq_epsilon)
  if hparams.bottleneck_kind == "em":
    return vq_discrete_bottleneck(
        x,
        hparams.bottleneck_bits,
        hparams.vq_beta,
        hparams.vq_decay,
        hparams.vq_epsilon,
        soft_em=True,
        num_samples=hparams.vq_num_samples)
  if hparams.bottleneck_kind == "gumbel_softmax":
    return gumbel_softmax_discrete_bottleneck(
        x,
        hparams.bottleneck_bits,
        hparams.vq_beta,
        hparams.vq_decay,
        hparams.vq_epsilon,
        hparams.temperature_warmup_steps,
        hard=False,
        summary=True)

  raise ValueError(
      "Unsupported hparams.bottleneck_kind %s" % hparams.bottleneck_kind)
1365593
def inference(self, kern_r, kern_c, Xr, Xc, Zr, Zc, likelihood, Y, qU_mean ,qU_var_r, qU_var_c):
        

        N, D, Mr, Mc, Qr, Qc = Y.shape[0], Y.shape[1], Zr.shape[0], Zc.shape[0], Zr.shape[1], Zc.shape[1]

        uncertain_inputs_r = isinstance(Xr, VariationalPosterior)
        uncertain_inputs_c = isinstance(Xc, VariationalPosterior)
        uncertain_outputs = isinstance(Y, VariationalPosterior)

        beta = 1./likelihood.variance

        psi0_r, psi1_r, psi2_r = self.gatherPsiStat(kern_r, Xr, Zr, uncertain_inputs_r)
        psi0_c, psi1_c, psi2_c = self.gatherPsiStat(kern_c, Xc, Zc, uncertain_inputs_c)

        #======================================================================
        # Compute Common Components
        #======================================================================

        Kuu_r = kern_r.K(Zr).copy()
        diag.add(Kuu_r, self.const_jitter)
        Lr = jitchol(Kuu_r)

        Kuu_c = kern_c.K(Zc).copy()
        diag.add(Kuu_c, self.const_jitter)
        Lc = jitchol(Kuu_c)

        mu, Sr, Sc = qU_mean, qU_var_r, qU_var_c
        LSr = jitchol(Sr)
        LSc = jitchol(Sc)

        LcInvMLrInvT = dtrtrs(Lc,dtrtrs(Lr,mu.T)[0].T)[0]
        LcInvPsi2_cLcInvT = backsub_both_sides(Lc, psi2_c,'right')
        LrInvPsi2_rLrInvT = backsub_both_sides(Lr, psi2_r,'right')
        LcInvLSc = dtrtrs(Lc, LSc)[0]
        LrInvLSr = dtrtrs(Lr, LSr)[0]
        LcInvScLcInvT = tdot(LcInvLSc)
        LrInvSrLrInvT = tdot(LrInvLSr)
        LcInvPsi1_cT = dtrtrs(Lc, psi1_c.T)[0]
        LrInvPsi1_rT = dtrtrs(Lr, psi1_r.T)[0]

        tr_LrInvPsi2_rLrInvT_LrInvSrLrInvT = (LrInvPsi2_rLrInvT*LrInvSrLrInvT).sum()
        tr_LcInvPsi2_cLcInvT_LcInvScLcInvT = (LcInvPsi2_cLcInvT*LcInvScLcInvT).sum()
        tr_LrInvSrLrInvT = np.square(LrInvLSr).sum()
        tr_LcInvScLcInvT = np.square(LcInvLSc).sum()
        tr_LrInvPsi2_rLrInvT = np.trace(LrInvPsi2_rLrInvT)
        tr_LcInvPsi2_cLcInvT = np.trace(LcInvPsi2_cLcInvT)

        #======================================================================
        # Compute log-likelihood
        #======================================================================

        logL_A = - np.square(Y).sum() \
               - (LcInvMLrInvT.T.dot(LcInvPsi2_cLcInvT).dot(LcInvMLrInvT)*LrInvPsi2_rLrInvT).sum() \
               -  tr_LrInvPsi2_rLrInvT_LrInvSrLrInvT* tr_LcInvPsi2_cLcInvT_LcInvScLcInvT \
               + 2 * (Y * LcInvPsi1_cT.T.dot(LcInvMLrInvT).dot(LrInvPsi1_rT)).sum() - psi0_c * psi0_r \
               + tr_LrInvPsi2_rLrInvT * tr_LcInvPsi2_cLcInvT

        logL = -N*D/2.*(np.log(2.*np.pi)-np.log(beta)) + beta/2.* logL_A \
               -Mc * (np.log(np.diag(Lr)).sum()-np.log(np.diag(LSr)).sum())  -Mr * (np.log(np.diag(Lc)).sum()-np.log(np.diag(LSc)).sum()) \
               - np.square(LcInvMLrInvT).sum()/2. - tr_LrInvSrLrInvT * tr_LcInvScLcInvT/2. + Mr*Mc/2.

        #======================================================================
        # Compute dL_dKuu
        #======================================================================

        tmp =  beta* LcInvPsi2_cLcInvT.dot(LcInvMLrInvT).dot(LrInvPsi2_rLrInvT).dot(LcInvMLrInvT.T) \
             + beta* tr_LrInvPsi2_rLrInvT_LrInvSrLrInvT * LcInvPsi2_cLcInvT.dot(LcInvScLcInvT) \
             - beta* LcInvMLrInvT.dot(LrInvPsi1_rT).dot(Y.T).dot(LcInvPsi1_cT.T) \
             - beta/2. * tr_LrInvPsi2_rLrInvT* LcInvPsi2_cLcInvT - Mr/2.*np.eye(Mc) \
             + tdot(LcInvMLrInvT)/2. + tr_LrInvSrLrInvT/2. * LcInvScLcInvT

        dL_dKuu_c = backsub_both_sides(Lc, tmp, 'left')
        dL_dKuu_c += dL_dKuu_c.T
        dL_dKuu_c *= 0.5

        tmp =  beta* LcInvMLrInvT.T.dot(LcInvPsi2_cLcInvT).dot(LcInvMLrInvT).dot(LrInvPsi2_rLrInvT) \
             + beta* tr_LcInvPsi2_cLcInvT_LcInvScLcInvT * LrInvPsi2_rLrInvT.dot(LrInvSrLrInvT) \
             - beta* LrInvPsi1_rT.dot(Y.T).dot(LcInvPsi1_cT.T).dot(LcInvMLrInvT) \
             - beta/2. * tr_LcInvPsi2_cLcInvT * LrInvPsi2_rLrInvT - Mc/2.*np.eye(Mr) \
             + tdot(LcInvMLrInvT.T)/2. + tr_LcInvScLcInvT/2. * LrInvSrLrInvT

        dL_dKuu_r = backsub_both_sides(Lr, tmp, 'left')
        dL_dKuu_r += dL_dKuu_r.T
        dL_dKuu_r *= 0.5

        #======================================================================
        # Compute dL_dthetaL
        #======================================================================

        dL_dthetaL = -D*N*beta/2. - logL_A*beta*beta/2.

        #======================================================================
        # Compute dL_dqU
        #======================================================================

        tmp = -beta * LcInvPsi2_cLcInvT.dot(LcInvMLrInvT).dot(LrInvPsi2_rLrInvT)\
              + beta* LcInvPsi1_cT.dot(Y).dot(LrInvPsi1_rT.T) - LcInvMLrInvT

        dL_dqU_mean = dtrtrs(Lc, dtrtrs(Lr, tmp.T, trans=1)[0].T, trans=1)[0]

        LScInv = dtrtri(LSc)
        tmp = -beta/2.*tr_LrInvPsi2_rLrInvT_LrInvSrLrInvT * LcInvPsi2_cLcInvT -tr_LrInvSrLrInvT/2.*np.eye(Mc)
        dL_dqU_var_c = backsub_both_sides(Lc, tmp, 'left') + tdot(LScInv.T) * Mr/2.

        LSrInv = dtrtri(LSr)
        tmp = -beta/2.*tr_LcInvPsi2_cLcInvT_LcInvScLcInvT * LrInvPsi2_rLrInvT -tr_LcInvScLcInvT/2.*np.eye(Mr)
        dL_dqU_var_r = backsub_both_sides(Lr, tmp, 'left') + tdot(LSrInv.T) * Mc/2.

        #======================================================================
        # Compute the Posterior distribution of inducing points p(u|Y)
        #======================================================================

        post = PosteriorMultioutput(LcInvMLrInvT=LcInvMLrInvT, LcInvScLcInvT=LcInvScLcInvT,
                LrInvSrLrInvT=LrInvSrLrInvT, Lr=Lr, Lc=Lc, kern_r=kern_r, Xr=Xr, Zr=Zr)

        #======================================================================
        # Compute dL_dpsi
        #======================================================================

        dL_dpsi0_r = - psi0_c * beta/2. * np.ones((D,))
        dL_dpsi0_c = - psi0_r * beta/2. * np.ones((N,))

        dL_dpsi1_c = beta * dtrtrs(Lc, (Y.dot(LrInvPsi1_rT.T).dot(LcInvMLrInvT.T)).T, trans=1)[0].T
        dL_dpsi1_r = beta * dtrtrs(Lr, (Y.T.dot(LcInvPsi1_cT.T).dot(LcInvMLrInvT)).T, trans=1)[0].T

        tmp = beta/2.*(-LcInvMLrInvT.dot(LrInvPsi2_rLrInvT).dot(LcInvMLrInvT.T) - tr_LrInvPsi2_rLrInvT_LrInvSrLrInvT * LcInvScLcInvT
              +tr_LrInvPsi2_rLrInvT *np.eye(Mc))
        dL_dpsi2_c = backsub_both_sides(Lc, tmp, 'left')
        tmp = beta/2.*(-LcInvMLrInvT.T.dot(LcInvPsi2_cLcInvT).dot(LcInvMLrInvT) - tr_LcInvPsi2_cLcInvT_LcInvScLcInvT * LrInvSrLrInvT
              +tr_LcInvPsi2_cLcInvT *np.eye(Mr))
        dL_dpsi2_r = backsub_both_sides(Lr, tmp, 'left')

        if not uncertain_inputs_r:
            dL_dpsi1_r += psi1_r.dot(dL_dpsi2_r+dL_dpsi2_r.T)
        if not uncertain_inputs_c:
            dL_dpsi1_c += psi1_c.dot(dL_dpsi2_c+dL_dpsi2_c.T)

        grad_dict = {
            'dL_dthetaL':dL_dthetaL,
            'dL_dqU_mean':dL_dqU_mean,
            'dL_dqU_var_c':dL_dqU_var_c,
            'dL_dqU_var_r':dL_dqU_var_r,
            'dL_dKuu_c': dL_dKuu_c,
            'dL_dKuu_r': dL_dKuu_r,
        }

        if uncertain_inputs_c:
            grad_dict['dL_dpsi0_c'] = dL_dpsi0_c
            grad_dict['dL_dpsi1_c'] = dL_dpsi1_c
            grad_dict['dL_dpsi2_c'] = dL_dpsi2_c
        else:
            grad_dict['dL_dKdiag_c'] = dL_dpsi0_c
            grad_dict['dL_dKfu_c'] = dL_dpsi1_c

        if uncertain_inputs_r:
            grad_dict['dL_dpsi0_r'] = dL_dpsi0_r
            grad_dict['dL_dpsi1_r'] = dL_dpsi1_r
            grad_dict['dL_dpsi2_r'] = dL_dpsi2_r
        else:
            grad_dict['dL_dKdiag_r'] = dL_dpsi0_r
            grad_dict['dL_dKfu_r'] = dL_dpsi1_r

        return post, logL, grad_dict
166985
def gene_expression_conv_base():
  """"""
  hparams = common_hparams.basic_params1()

  batch_size = 10
  output_length = 2048
  inputs_per_output = 128
  chunk_size = 4
  input_length = output_length * inputs_per_output // chunk_size
  hparams.batch_size = input_length * batch_size

  hparams.dropout = 0.1
  hparams.add_hparam("num_conv_layers", 4)
  hparams.add_hparam("num_dconv_layers", 7)
  # The product of these pooling windows should match
  # input_length/target_length.
  hparams.add_hparam("pooling_windows", [2, 2, 2, 4])

  hparams.hidden_size = 256
  hparams.kernel_width = 20
  hparams.add_hparam("stride", 1)
  return hparams
69253
def transformer_wikitext103_l4k_v0():
  """"""
  hparams = transformer_big()

  # Adafactor uses less memory than Adam.
  # switch to Adafactor with its recommended learning rate scheme.
  hparams.optimizer = "Adafactor"
  hparams.learning_rate_schedule = "rsqrt_decay"
  hparams.learning_rate_warmup_steps = 10000

  hparams.num_heads = 4
  hparams.max_length = 4096
  hparams.batch_size = 4096
  hparams.shared_embedding_and_softmax_weights = False

  hparams.num_hidden_layers = 8
  hparams.attention_dropout = 0.1
  hparams.layer_prepostprocess_dropout = 0.2
  hparams.relu_dropout = 0.1
  hparams.label_smoothing = 0.0

  # Using noise broadcast in the dropout layers saves memory during training.
  hparams.attention_dropout_broadcast_dims = "0,1"  # batch, heads
  hparams.relu_dropout_broadcast_dims = "1"  # length
  hparams.layer_prepostprocess_dropout_broadcast_dims = "1"  # length

  # Avoid an expensive concat on TPU.
  # >1 shards helps with faster parameter distribution on multi-GPU machines
  hparams.symbol_modality_num_shards = 1

  return hparams
880520
def multi_conv_res(x, padding, name, layers, hparams, mask=None, source=None):
  """"""
  with tf.variable_scope(name):
    padding_bias = None
    if mask is not None:
      padding_bias = (1.0 - mask) * -1e9  # Bias to not attend to padding.
      if padding == "LEFT":  # Do not mask anything when left-padding.
        mask = None
    if (hparams.kernel_scheme in _KERNEL_SCHEMES and
        hparams.dilation_scheme in _DILATION_SCHEMES):
      kernels = _KERNEL_SCHEMES[hparams.kernel_scheme]
      dilations = _DILATION_SCHEMES[hparams.dilation_scheme]
      dilations_and_kernels = list(zip(dilations, kernels))
      dilations_and_kernels1 = dilations_and_kernels[:2]
      dilations_and_kernels2 = dilations_and_kernels[2:]
    else:
      k = (hparams.kernel_height, hparams.kernel_width)
      k2 = (hparams.large_kernel_size, 1)
      dilations_and_kernels1 = [((1, 1), k), ((1, 1), k)]
      dilations_and_kernels2 = [((1, 1), k2), ((4, 4), k2)]
    separabilities1 = [hparams.separability, hparams.separability]
    separabilities2 = [hparams.separability] * len(dilations_and_kernels2)
    if hparams.separability < 0:
      separabilities1 = [hparams.separability - 1, hparams.separability]
      separabilities2 = [
          hparams.separability - i
          for i in reversed(range(len(dilations_and_kernels2)))
      ]

    def norm_fn(x, name):
      with tf.variable_scope(name, default_name="norm"):
        return common_layers.apply_norm(
            x, hparams.norm_type, hparams.hidden_size, hparams.norm_epsilon)

    for layer in range(layers):
      with tf.variable_scope("layer_%d" % layer):
        y = common_layers.subseparable_conv_block(
            x,
            hparams.hidden_size,
            dilations_and_kernels1,
            normalizer_fn=norm_fn,
            padding=padding,
            mask=mask,
            separabilities=separabilities1,
            name="residual1")
        x += common_layers.subseparable_conv_block(
            x + y,
            hparams.hidden_size,
            dilations_and_kernels2,
            normalizer_fn=norm_fn,
            padding=padding,
            mask=mask,
            separabilities=separabilities2,
            name="residual2") + y
        if source is not None and hparams.attention_type != "none":
          x += attention(x, source, norm_fn, hparams, bias=padding_bias)
        if mask is not None:
          x *= mask
    return tf.nn.dropout(x, 1.0 - hparams.dropout)
1187761
def imagetransformerpp_sep_channels_8l_8h():
  """"""
  hparams = imagetransformer_base()
  hparams.likelihood = cia.DistributionType.DMOL
  hparams.num_channels = 1
  hparams.bottom["targets"] = modalities.image_channel_compress_targets_bottom
  hparams.top["targets"] = modalities.identity_top
  hparams.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
188936
def fitlin(imgarr,refarr):
    
    # Initialize variables
    _mat = np.zeros((3,3),dtype=np.float64)
    _xorg = imgarr[0][0]
    _yorg = imgarr[0][1]
    _xoorg = refarr[0][0]
    _yoorg = refarr[0][1]
    _sigxox = 0.
    _sigxoy = 0.
    _sigxo = 0.
    _sigyox = 0.
    _sigyoy = 0.
    _sigyo = 0.

    _npos = len(imgarr)
    # Populate matrices
    for i in range(_npos):
        _mat[0][0] += np.power((imgarr[i][0] - _xorg),2)
        _mat[0][1] += (imgarr[i][0] - _xorg) * (imgarr[i][1] - _yorg)
        _mat[0][2] += (imgarr[i][0] - _xorg)
        _mat[1][1] += np.power((imgarr[i][1] - _yorg),2)
        _mat[1][2] += imgarr[i][1] - _yorg

        _sigxox += (refarr[i][0] - _xoorg)*(imgarr[i][0] - _xorg)
        _sigxoy += (refarr[i][0] - _xoorg)*(imgarr[i][1] - _yorg)
        _sigxo += refarr[i][0] - _xoorg
        _sigyox += (refarr[i][1] - _yoorg)*(imgarr[i][0] -_xorg)
        _sigyoy += (refarr[i][1] - _yoorg)*(imgarr[i][1] - _yorg)
        _sigyo += refarr[i][1] - _yoorg

    _mat[2][2] = _npos
    _mat[1][0] = _mat[0][1]
    _mat[2][0] = _mat[0][2]
    _mat[2][1] = _mat[1][2]

    # Now invert this matrix
    _mat = linalg.inv(_mat)

    _a  = _sigxox*_mat[0][0]+_sigxoy*_mat[0][1]+_sigxo*_mat[0][2]
    _b  = -1*(_sigxox*_mat[1][0]+_sigxoy*_mat[1][1]+_sigxo*_mat[1][2])
    #_x0 = _sigxox*_mat[2][0]+_sigxoy*_mat[2][1]+_sigxo*_mat[2][2]

    _c  = _sigyox*_mat[1][0]+_sigyoy*_mat[1][1]+_sigyo*_mat[1][2]
    _d  = _sigyox*_mat[0][0]+_sigyoy*_mat[0][1]+_sigyo*_mat[0][2]
    #_y0 = _sigyox*_mat[2][0]+_sigyoy*_mat[2][1]+_sigyo*_mat[2][2]

    _xt = _xoorg - _a*_xorg+_b*_yorg
    _yt = _yoorg - _d*_xorg-_c*_yorg

    return [_a,_b,_xt],[_c,_d,_yt]
885099
def lmx_base():
  """"""
  hparams = transformer.transformer_tpu()
  # sharing is counterproductive when underparameterized
  hparams.shared_embedding_and_softmax_weights = False
  # we judge by log-ppl, so label smoothing hurts.
  hparams.label_smoothing = 0.0
  # This makes the batch size on GPU the same as on TPU for a packed problem
  # with sequence length 256.
  # TODO(noam): fix the mess that is the data reading pipeline.
  hparams.max_length = 256
  # larger batch since we only have a decoder
  hparams.batch_size = 4096
  # save some memory so we can have a larger model
  hparams.activation_dtype = "bfloat16"
  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"])
449152
def multifuncrs(X, funs, eps=1E-6,
                nswp=10,
                kickrank=5,
                y0=None,
                rmax=999999,  # TODO:infinity \
                kicktype='amr-two',        \
                pcatype='svd',             \
                trunctype='fro',           \
                d2=1,                      \
                do_qr=False,               \
                verb=1):
    

    dtype = np.float64
    if len([x for x in X if x.is_complex]) > 0:
        dtype = np.complex128

    y = y0
    wasrand = False

    nx = len(X)
    d = X[0].d
    n = X[0].n
    rx = np.transpose(np.array([ttx.r for ttx in X]))
    #crx = [tt.tensor.to_list(ttx) for x in X]
    #crx = zip(*crx)
    crx = np.transpose(np.array([tt.tensor.to_list(ttx)
                                 for ttx in X], dtype=np.object))
    crx = np.empty((nx, d), dtype=np.object)
    i = 0
    for ttx in X:
        v = tt.tensor.to_list(ttx)
        j = 0
        for w in v:
            crx[i, j] = w
            j = j + 1
        i = i + 1
    crx = crx.T
    if y is None:
        ry = d2 * np.ones((d + 1,), dtype=np.int32)
        ry[0] = 1
        y = tt.rand(n, d, ry)
        wasrand = True

    ry = y.r
    cry = tt.tensor.to_list(y)

    Ry = np.zeros((d + 1, ), dtype=np.object)
    Ry[0] = np.array([[1.0]], dtype=dtype)
    Ry[d] = np.array([[1.0]], dtype=dtype)
    Rx = np.zeros((d + 1, nx), dtype=np.object)
    Rx[0, :] = np.ones(nx, dtype=dtype)
    Rx[d, :] = np.ones(nx, dtype=dtype)

    block_order = [+d, -d]

    # orth
    for i in range(0, d - 1):
        cr = cry[i]
        cr = reshape(cr, (ry[i] * n[i], ry[i + 1]))
        cr, rv = np.linalg.qr(cr)
        cr2 = cry[i + 1]
        cr2 = reshape(cr2, (ry[i + 1], n[i + 1] * ry[i + 2]))
        cr2 = np.dot(rv, cr2)  # matrix multiplication
        ry[i + 1] = cr.shape[1]
        cr = reshape(cr, (ry[i], n[i], ry[i + 1]))
        cry[i + 1] = reshape(cr2, (ry[i + 1], n[i + 1], ry[i + 2]))
        cry[i] = cr

        Ry[i + 1] = np.dot(Ry[i], reshape(cr, (ry[i], n[i] * ry[i + 1])))
        Ry[i + 1] = reshape(Ry[i + 1], (ry[i] * n[i], ry[i + 1]))
        curind = []
        if wasrand:
            # EVERY DAY I'M SHUFFLIN'
            curind = np.random.permutation(n[i] * ry[i])[:ry[i + 1]]
        else:
            curind = maxvol(Ry[i + 1])
        Ry[i + 1] = Ry[i + 1][curind, :]
        for j in range(0, nx):
            try:
                Rx[i + 1, j] = reshape(crx[i, j],
                                       (rx[i, j], n[i] * rx[i + 1, j]))
            except:
                pass
            Rx[i + 1, j] = np.dot(Rx[i, j], Rx[i + 1, j])
            Rx[i + 1, j] = reshape(Rx[i + 1, j], (ry[i] * n[i], rx[i + 1, j]))
            Rx[i + 1, j] = Rx[i + 1, j][curind, :]

    d2 = ry[d]
    ry[d] = 1
    cry[d - 1] = np.transpose(cry[d - 1], [2, 0, 1])  # permute

    last_sweep = False
    swp = 1

    dy = np.zeros((d, ))
    max_dy = 0

    cur_order = copy.copy(block_order)
    order_index = 1
    i = d - 1
    # can't use 'dir' identifier in python
    dirn = int(math.copysign(1, cur_order[order_index]))

    # DMRG sweeps
    while swp <= nswp or dirn > 0:

        oldy = reshape(cry[i], (d2 * ry[i] * n[i] * ry[i + 1],))

        if not last_sweep:
            # compute the X superblocks
            curbl = np.zeros((ry[i] * n[i] * ry[i + 1], nx), dtype=dtype)
            for j in range(0, nx):
                cr = reshape(crx[i, j], (rx[i, j], n[i] * rx[i + 1, j]))
                cr = np.dot(Rx[i, j], cr)
                cr = reshape(cr, (ry[i] * n[i], rx[i + 1, j]))
                cr = np.dot(cr, Rx[i + 1, j])
                curbl[:, j] = cr.flatten('F')
            # call the function
            newy = funs(curbl)
            # multiply with inverted Ry
            newy = reshape(newy, (ry[i], n[i] * ry[i + 1] * d2))
            newy = np.linalg.solve(Ry[i], newy)  # y = R \ y
            newy = reshape(newy, (ry[i] * n[i] * ry[i + 1], d2))
            newy = reshape(np.transpose(newy), (d2 * ry[i] * n[i], ry[i + 1]))
            newy = np.transpose(np.linalg.solve(
                np.transpose(Ry[i + 1]), np.transpose(newy)))  # y=y/R
            newy = reshape(newy, (d2 * ry[i] * n[i] * ry[i + 1],))
        else:
            newy = oldy

        dy[i] = np.linalg.norm(newy - oldy) / np.linalg.norm(newy)
        max_dy = max(max_dy, dy[i])

        # truncation
        if dirn > 0:  # left-to-right
            newy = reshape(newy, (d2, ry[i] * n[i] * ry[i + 1]))
            newy = reshape(np.transpose(newy), (ry[i] * n[i], ry[i + 1] * d2))
        else:
            newy = reshape(newy, (d2 * ry[i], n[i] * ry[i + 1]))

        r = 0  # defines a variable in global scope

        if kickrank >= 0:
            u, s, v = np.linalg.svd(newy, full_matrices=False)
            v = np.conj(np.transpose(v))
            if trunctype == "fro" or last_sweep:
                r = my_chop2(s, eps / math.sqrt(d) * np.linalg.norm(s))
            else:
                # truncate taking into account the (r+1) overhead in the cross
                # (T.S.: what?)
                cums = abs(s * np.arange(2, len(s) + 2)) ** 2
                cums = np.cumsum(cums[::-1])[::-1]
                cums = cums / cums[0]
                ff = [i for i in range(len(cums)) if cums[i] < eps ** 2 / d]
                if len(ff) == 0:
                    r = len(s)
                else:
                    r = np.amin(ff)
            r = min(r, rmax, len(s))
        else:
            if dirn > 0:
                u, v = np.linalg.qr(newy)
                v = np.conj(np.transpose(v))
                r = u.shape[1]
                s = np.ones((r, ))
            else:
                v, u = np.linalg.qr(np.transpose(newy))
                v = np.conj(v)
                u = np.transpose(u)
                r = u.shape[1]
                s = np.ones((r, ))

        if verb > 1:
            print('=multifuncrs=   block %d{%d}, dy: %3.3e, r: %d' % (i, dirn, dy[i], r))

        # kicks and interfaces
        if dirn > 0 and i < d - 1:
            u = u[:, :r]
            v = np.dot(v[:, :r], np.diag(s[:r]))

            # kick
            radd = 0
            rv = 1
            if not last_sweep and kickrank > 0:
                uk = None
                if kicktype == 'amr-two':
                    # AMR(two)-like kick.

                    # compute the X superblocks
                    ind2 = np.unique(np.random.randint(
                        0, ry[i + 2] * n[i + 1], ry[i + 1]))
                    #ind2 = np.unique(np.floor(np.random.rand(ry[i + 1]) * (ry[i + 2] * n[i + 1])))
                    rkick = len(ind2)
                    curbl = np.zeros((ry[i] * n[i] * rkick, nx), dtype=dtype)
                    for j in range(nx):
                        cr1 = reshape(
                            crx[i, j], (rx[i, j], n[i] * rx[i + 1, j]))
                        cr1 = np.dot(Rx[i, j], cr1)
                        cr1 = reshape(cr1, (ry[i] * n[i], rx[i + 1, j]))
                        cr2 = reshape(
                            crx[i + 1, j], (rx[i + 1, j] * n[i + 1], rx[i + 2, j]))
                        cr2 = np.dot(cr2, Rx[i + 2, j])
                        cr2 = reshape(
                            cr2, (rx[i + 1, j], n[i + 1] * ry[i + 2]))
                        cr2 = cr2[:, ind2]
                        curbl[:, j] = reshape(
                            np.dot(cr1, cr2), (ry[i] * n[i] * rkick,))
                    # call the function
                    uk = funs(curbl)
                    uk = reshape(uk, (ry[i], n[i] * rkick * d2))
                    uk = np.linalg.solve(Ry[i], uk)
                    uk = reshape(uk, (ry[i] * n[i], rkick * d2))
                    if pcatype == 'svd':
                        uk, sk, vk = np.linalg.svd(uk, full_matrices=False)
                        vk = np.conj(np.transpose(vk))
                        uk = uk[:, :min(kickrank, uk.shape[1])]
                    else:
                        # uk = uchol(np.transpose(uk), kickrank + 1) # TODO
                        uk = uk[:, :max(uk.shape[1] - kickrank + 1, 1):-1]
                else:
                    uk = np.random.rand(ry[i] * n[i], kickrank)
                u, rv = np.linalg.qr(np.concatenate((u, uk), axis=1))
                radd = uk.shape[1]
            v = np.concatenate(
                (v, np.zeros((ry[i + 1] * d2, radd), dtype=dtype)), axis=1)
            v = np.dot(rv, np.conj(np.transpose(v)))
            r = u.shape[1]

            cr2 = cry[i + 1]
            cr2 = reshape(cr2, (ry[i + 1], n[i + 1] * ry[i + 2]))
            v = reshape(v, (r * ry[i + 1], d2))
            v = reshape(np.transpose(v), (d2 * r, ry[i + 1]))
            v = np.dot(v, cr2)

            ry[i + 1] = r

            u = reshape(u, (ry[i], n[i], r))
            v = reshape(v, (d2, r, n[i + 1], ry[i + 2]))

            cry[i] = u
            cry[i + 1] = v

            Ry[i + 1] = np.dot(Ry[i], reshape(u, (ry[i], n[i] * ry[i + 1])))
            Ry[i + 1] = reshape(Ry[i + 1], (ry[i] * n[i], ry[i + 1]))
            curind = maxvol(Ry[i + 1])
            Ry[i + 1] = Ry[i + 1][curind, :]
            for j in range(nx):
                Rx[i + 1, j] = reshape(crx[i, j],
                                       (rx[i, j], n[i] * rx[i + 1, j]))
                Rx[i + 1, j] = np.dot(Rx[i, j], Rx[i + 1, j])
                Rx[i + 1, j] = reshape(Rx[i + 1, j],
                                       (ry[i] * n[i], rx[i + 1, j]))
                Rx[i + 1, j] = Rx[i + 1, j][curind, :]
        elif dirn < 0 and i > 0:
            u = np.dot(u[:, :r], np.diag(s[:r]))
            v = np.conj(v[:, :r])

            radd = 0
            rv = 1
            if not last_sweep and kickrank > 0:
                if kicktype == 'amr-two':
                    # compute the X superblocks
                    ind2 = np.unique(np.random.randint(
                        0, ry[i - 1] * n[i - 1], ry[i]))
                    rkick = len(ind2)
                    curbl = np.zeros(
                        (rkick * n[i] * ry[i + 1], nx), dtype=dtype)
                    for j in range(nx):
                        cr1 = reshape(
                            crx[i, j], (rx[i, j] * n[i], rx[i + 1, j]))
                        cr1 = np.dot(cr1, Rx[i + 1, j])
                        cr1 = reshape(cr1, (rx[i, j], n[i] * ry[i + 1]))
                        cr2 = reshape(
                            crx[i - 1, j], (rx[i - 1, j], n[i - 1] * rx[i, j]))
                        cr2 = np.dot(Rx[i - 1, j], cr2)
                        cr2 = reshape(cr2, (ry[i - 1] * n[i - 1], rx[i, j]))
                        cr2 = cr2[ind2, :]
                        curbl[:, j] = reshape(
                            np.dot(cr2, cr1), (rkick * n[i] * ry[i + 1],))
                    # calling the function
                    uk = funs(curbl)
                    uk = reshape(uk, (rkick * n[i] * ry[i + 1], d2))
                    uk = reshape(np.transpose(
                        uk), (d2 * rkick * n[i], ry[i + 1]))
                    uk = np.transpose(np.linalg.solve(
                        np.transpose(Ry[i + 1]), np.transpose(uk)))
                    uk = reshape(uk, (d2 * rkick, n[i] * ry[i + 1]))
                    if pcatype == 'svd':
                        vk, sk, uk = np.linalg.svd(uk, full_matrices=False)
                        uk = np.conj(np.transpose(uk))
                        # TODO: refactor
                        uk = uk[:, :min(kickrank, uk.shape[1])]
                    else:
                        # uk = uchol(uk, kickrank + 1) # TODO
                        uk = uk[:, :max(uk.shape[1] - kickrank + 1, 1):-1]
                else:
                    uk = np.random.rand(n[i] * ry[i + 1], kickrank)
                v, rv = np.linalg.qr(np.concatenate((v, uk), axis=1))
                radd = uk.shape[1]
            u = np.concatenate(
                (u, np.zeros((d2 * ry[i], radd), dtype=dtype)), axis=1)
            u = np.dot(u, np.transpose(rv))
            r = v.shape[1]
            cr2 = cry[i - 1]
            cr2 = reshape(cr2, (ry[i - 1] * n[i - 1], ry[i]))
            u = reshape(u, (d2, ry[i] * r))
            u = reshape(np.transpose(u), (ry[i], r * d2))
            u = np.dot(cr2, u)

            u = reshape(u, (ry[i - 1] * n[i - 1] * r, d2))
            u = reshape(np.transpose(u), (d2, ry[i - 1], n[i - 1], r))
            v = reshape(np.transpose(v), (r, n[i], ry[i + 1]))

            ry[i] = r
            cry[i - 1] = u
            cry[i] = v

            Ry[i] = np.dot(reshape(v, (ry[i] * n[i], ry[i + 1])), Ry[i + 1])
            Ry[i] = reshape(Ry[i], (ry[i], n[i] * ry[i + 1]))
            curind = maxvol(np.transpose(Ry[i]))
            Ry[i] = Ry[i][:, curind]
            for j in range(nx):
                Rx[i, j] = reshape(crx[i, j], (rx[i, j] * n[i], rx[i + 1, j]))
                Rx[i, j] = np.dot(Rx[i, j], Rx[i + 1, j])
                Rx[i, j] = reshape(Rx[i, j], (rx[i, j], n[i] * ry[i + 1]))
                Rx[i, j] = Rx[i, j][:, curind]
        elif dirn > 0 and i == d - 1:
            newy = np.dot(np.dot(u[:, :r], np.diag(s[:r])),
                          np.conj(np.transpose(v[:, :r])))
            newy = reshape(newy, (ry[i] * n[i] * ry[i + 1], d2))
            cry[i] = reshape(np.transpose(newy), (d2, ry[i], n[i], ry[i + 1]))
        elif dirn < 0 and i == 0:
            newy = np.dot(np.dot(u[:, :r], np.diag(s[:r])),
                          np.conj(np.transpose(v[:, :r])))
            newy = reshape(newy, (d2, ry[i], n[i], ry[i + 1]))
            cry[i] = newy

        i = i + dirn
        cur_order[order_index] = cur_order[order_index] - dirn
        if cur_order[order_index] == 0:
            order_index = order_index + 1
            if verb > 0:
                print('=multifuncrs= sweep %d{%d}, max_dy: %3.3e, erank: %g' % (swp, order_index, max_dy,
                                                                                math.sqrt(np.dot(ry[:d], n * ry[1:]) / np.sum(n))))

            if last_sweep:
                break
            if max_dy < eps and dirn < 0:
                last_sweep = True
                kickrank = 0

            if order_index >= len(cur_order):
                cur_order = copy.copy(block_order)
                order_index = 0
                if last_sweep:
                    cur_order = [d - 1]

                max_dy = 0
                swp = swp + 1

            dirn = int(math.copysign(1, cur_order[order_index]))
            i = i + dirn

    cry[d - 1] = np.transpose(cry[d - 1][:, :, :, 0], [1, 2, 0])
    y = tt.tensor.from_list(cry)
    return y
187739
def autoencoder_residual_discrete_big():
  """"""
  hparams = autoencoder_residual_discrete()
  hparams.hidden_size = 128
  hparams.max_hidden_size = 4096
  hparams.bottleneck_noise = 0.1
  hparams.residual_dropout = 0.4
  return hparams
220287
def mtf_transformer_base():
  """"""
  hparams = common_hparams.basic_params1()
  hparams.no_data_parallelism = True
  hparams.use_fixed_batch_size = True
  hparams.add_hparam("mtf_mode", True)
  hparams.batch_size = 64
  hparams.max_length = 256
  hparams.add_hparam("d_model", 512)
  hparams.add_hparam("d_kv", 128)
  hparams.add_hparam("local_attention_window_size", 128)
  hparams.label_smoothing = 0.1
  # 8-way model-parallelism
  hparams.add_hparam("mesh_shape", "model:8")
  hparams.add_hparam("layout", "batch:batch;vocab:model;d_ff:model;heads:model")
  hparams.add_hparam("num_heads", 8)
  hparams.add_hparam("d_ff", 2048)
  hparams.add_hparam("encoder_replicate_factor", 1)
  hparams.add_hparam("decoder_replicate_factor", 1)
  hparams.add_hparam("encoder_layers", ["att", "drd"] * 6)
  hparams.add_hparam("decoder_layers", ["att", "enc_att", "drd"] * 6)
  hparams.add_hparam("attention_dropout", 0.1)
  hparams.add_hparam("relu_dropout", 0.1)
  hparams.layer_prepostprocess_dropout = 0.1

  # Describes what model architecture:
  #   "encdec": encoder + autoregressive decoder
  #   "decoder": single-stack autoregressive sequence model.
  #   "encoder": single-stack non-autoregressive model
  #      with equal-length inputs and outputs.
  hparams.add_hparam("transformer_type", "encdec")

  # What does the decoder do:
  #   "autoregressive": Decoder left to right
  #   "denoising": Fills in masked-out values simultaneously
  hparams.add_hparam("decoder_type", "autoregressive")

  # Parameters describing the noising algorithm for denoising decoders
  hparams.add_hparam("noising_spec_train", {"type": "mask", "prob": 0.15})
  hparams.add_hparam("noising_spec_eval", {"type": "mask", "prob": 0.15})
  # during training, we use the eval noiser with this probability
  hparams.add_hparam("noising_use_eval_during_train", 0.1)

  # round up vocab sizes to be a multiple of this value
  hparams.vocab_divisor = 128

  # options are dense_relu_dense, moe, hmoe
  hparams.add_hparam("feedforward_layer", "drd")

  # If True, then reuse targets_embedding_var * rsqrt(d_model) as softmax_var
  # If hparams.transformer_type == "encoder", then there is no targets embedding
  # so we reuse the inputs embedding instead.
  hparams.shared_embedding_and_softmax_weights = True
  # Reuse targets_embedding_var as inputs_embedding_var
  # relevant only if hparams.transformer_type == "encdec"
  hparams.shared_embedding = True
  hparams.optimizer = "Adafactor"
  hparams.learning_rate_schedule = "linear_warmup*rsqrt_decay*linear_decay"
  hparams.learning_rate_warmup_steps = 10000
  hparams.add_hparam("master_dtype", "bfloat16")
  hparams.add_hparam("slice_dtype", "float32")
  hparams.activation_dtype = "bfloat16"

  # These parameters make Transformer model compatible with MtfTransformer
  # Do not override these, as mtf_transformer does not support other options.
  hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
  hparams.bottom = {
      "inputs": modalities.identity_bottom,
      "targets": modalities.identity_bottom,
  }
  hparams.top = {
      "targets": modalities.identity_top,
  }

  # Parameters for computing the maximum decode length in beam search.
  # Maximum decode length is:
  #    min(max_length,
  #        decode_length_multiplier * input_length + decode_length_constant)
  hparams.add_hparam("decode_length_multiplier", 1.5)
  hparams.add_hparam("decode_length_constant", 10.0)

  # If nonzero, we split the batch across two tensor-dimensions named
  # "outer_batch" and "inner_batch", allowing for splitting across two mesh
  # dimensions.  This is necessary for hierarchical mixture of experts.
  # The two tensor dimensions have sizes hparams.outer_batch_size and
  # hparams.batch_size // hparams.outer_batch_size.
  hparams.add_hparam("outer_batch_size", 0)

  # TODO(noam): file a bug
  hparams.add_hparam("reshape_logits_hack", False)
  hparams.add_hparam("compression_factor", 4)

  return hparams
1130838
def transformer_imagenet64_memory_v0():
  """"""
  hparams = transformer_cifar10_memory_v0()

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

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

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

  hparams.max_relative_position = 3072

  return hparams
368015
def optimize_rfc(data, targets):
    """"""
    def rfc_crossval(n_estimators, min_samples_split, max_features):
        """Wrapper of RandomForest cross validation.

        Notice how we ensure n_estimators and min_samples_split are casted
        to integer before we pass them along. Moreover, to avoid max_features
        taking values outside the (0, 1) range, we also ensure it is capped
        accordingly.
        """
        return rfc_cv(
            n_estimators=int(n_estimators),
            min_samples_split=int(min_samples_split),
            max_features=max(min(max_features, 0.999), 1e-3),
            data=data,
            targets=targets,
        )

    optimizer = BayesianOptimization(
        f=rfc_crossval,
        pbounds={
            "n_estimators": (10, 250),
            "min_samples_split": (2, 25),
            "max_features": (0.1, 0.999),
        },
        random_state=1234,
        verbose=2
    )
    optimizer.maximize(n_iter=10)

    print("Final result:", optimizer.max)
170789
def imagetransformer_b12l_4h_b128_h512_uncond_dr01_im():
  """"""
  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