Simple discretization through tanh, flip bottleneck_noise many bits.
186619
def tanh_discrete_bottleneck(x, bottleneck_bits, bottleneck_noise,
                             discretize_warmup_steps, mode):
  """"""
  x = tf.layers.dense(x, bottleneck_bits, name="tanh_discrete_bottleneck")
  d0 = tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x))) - 1.0
  if mode == tf.estimator.ModeKeys.TRAIN:
    x += tf.truncated_normal(
        common_layers.shape_list(x), mean=0.0, stddev=0.2)
  x = tf.tanh(x)
  d = x + tf.stop_gradient(2.0 * tf.to_float(tf.less(0.0, x)) - 1.0 - x)
  if mode == tf.estimator.ModeKeys.TRAIN:
    noise = tf.random_uniform(common_layers.shape_list(x))
    noise = 2.0 * tf.to_float(tf.less(bottleneck_noise, noise)) - 1.0
    d *= noise
  d = common_layers.mix(d, x, discretize_warmup_steps,
                        mode == tf.estimator.ModeKeys.TRAIN)
  return d, d0
936605
def traptransit_MCMC(ts,fs,dfs=1e-5,nwalkers=200,nburn=300,niter=1000,
                     threads=1,p0=[0.1,0.1,3,0],return_sampler=False,
                     maxslope=MAXSLOPE):
    
    model = TraptransitModel(ts,fs,dfs,maxslope=maxslope)
    sampler = emcee.EnsembleSampler(nwalkers,4,model,threads=threads)
    T0 = p0[0]*(1+rand.normal(size=nwalkers)*0.1)
    d0 = p0[1]*(1+rand.normal(size=nwalkers)*0.1)
    slope0 = p0[2]*(1+rand.normal(size=nwalkers)*0.1)
    ep0 = p0[3]+rand.normal(size=nwalkers)*0.0001

    p0 = np.array([T0,d0,slope0,ep0]).T

    pos, prob, state = sampler.run_mcmc(p0, nburn)
    sampler.reset()
    sampler.run_mcmc(pos, niter, rstate0=state)
    if return_sampler:
        return sampler
    else:
        return sampler.flatchain[:,0],sampler.flatchain[:,1],sampler.flatchain[:,2],sampler.flatchain[:,3]
54959
def susvd(x, x_obs, rho, penalties):
    

    assert type(x) == Tensor, "Input array must be a Tensor"

    while True:

        # proximal operator for the Fro. norm
        x = squared_error(x, rho, x_obs)

        # sequential singular value thresholding
        for ix, penalty in enumerate(penalties):
            x = x.unfold(ix).svt(penalty / rho).fold()

        yield x
1049749
def autoencoder_basic_discrete():
  """"""
  hparams = autoencoder_autoregressive()
  hparams.num_hidden_layers = 5
  hparams.hidden_size = 64
  hparams.bottleneck_bits = 1024
  hparams.bottleneck_noise = 0.1
  hparams.add_hparam("discretize_warmup_steps", 16000)
  return hparams
1245954
def conv2d(ni:int, nf:int, ks:int=3, stride:int=1, padding:int=None, bias=False, init:LayerFunc=nn.init.kaiming_normal_) -> nn.Conv2d:
    ""
    if padding is None: padding = ks//2
    return init_default(nn.Conv2d(ni, nf, kernel_size=ks, stride=stride, padding=padding, bias=bias), init)
23417
def mtf_image_transformer_tiny():
  """"""
  hparams = mtf_image_transformer_base()
  hparams.hidden_size = 128
  hparams.d_ff = 256
  hparams.batch_size = 4
  hparams.num_encoder_layers = 1
  hparams.num_decoder_layers = 4
  hparams.num_heads = 4
  hparams.attention_key_size = 128
  hparams.attention_value_size = 128
  hparams.block_length = 32
  # data parallelism and model-parallelism
  hparams.mesh_shape = "batch:2"
  hparams.layout = "batch:batch"
  return hparams
445569
def node_to_evenly_discretized(node):
    
    if not all([node.attrib["minMag"], node.attrib["binWidth"],
                node.nodes[0].text]):
        return None
    # Text to float
    rates = [float(x) for x in node.nodes[0].text.split()]
    return mfd.evenly_discretized.EvenlyDiscretizedMFD(
        float(node.attrib["minMag"]),
        float(node.attrib["binWidth"]),
        rates)
472808
def tanh_discrete_unbottleneck(x, hidden_size):
  """"""
  x = tf.layers.dense(x, hidden_size, name="tanh_discrete_unbottleneck")
  return x
1040257
def next_frame_basic_stochastic_discrete():
  """"""
  hparams = basic_deterministic_params.next_frame_sampling()
  hparams.batch_size = 4
  hparams.video_num_target_frames = 6
  hparams.scheduled_sampling_mode = "prob_inverse_lin"
  hparams.scheduled_sampling_decay_steps = 40000
  hparams.scheduled_sampling_max_prob = 1.0
  hparams.dropout = 0.15
  hparams.filter_double_steps = 3
  hparams.hidden_size = 96
  hparams.learning_rate_constant = 0.002
  hparams.learning_rate_warmup_steps = 2000
  hparams.learning_rate_schedule = "linear_warmup * constant"
  hparams.concat_internal_states = True
  hparams.video_modality_loss_cutoff = 0.03
  hparams.add_hparam("bottleneck_bits", 128)
  hparams.add_hparam("bottleneck_noise", 0.1)
  hparams.add_hparam("discretize_warmup_steps", 40000)
  hparams.add_hparam("latent_rnn_warmup_steps", 40000)
  hparams.add_hparam("latent_rnn_max_sampling", 0.5)
  hparams.add_hparam("latent_use_max_probability", 0.8)
  hparams.add_hparam("full_latent_tower", False)
  hparams.add_hparam("latent_predictor_state_size", 128)
  hparams.add_hparam("latent_predictor_temperature", 1.0)
  hparams.add_hparam("complex_addn", True)
  hparams.add_hparam("recurrent_state_size", 64)
  return hparams
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
421994
def next_frame_sv2p_discrete():
  """"""
  hparams = next_frame_sv2p()
  hparams.action_injection = "multiplicative"
  hparams.small_mode = True
  hparams.add_hparam("bottleneck_bits", 128)
  hparams.add_hparam("bottleneck_noise", 0.02)
  hparams.add_hparam("discrete_warmup_steps", 40000)
  hparams.add_hparam("full_latent_tower", False)
  hparams.add_hparam("latent_predictor_state_size", 128)
  hparams.add_hparam("latent_predictor_temperature", 0.5)
  hparams.add_hparam("discretize_warmup_steps", 40000)
  return hparams
23497
def _smooth(aprob, bprob, Ndiff):
    
    gamma, epsilon = _smoothing_parameters(aprob, bprob, Ndiff)

    # Remove zeros.
    in_a = [i for i,v in enumerate(aprob) if abs(v) > 0.0]
    aprob = list([list(aprob)[i] for i in in_a])
    bprob = list([list(bprob)[i]*gamma for i in in_a])

    # Replace zero values with epsilon.
    bprob = list(map(lambda v: v if v != 0. else epsilon, bprob))

    return aprob, bprob
722209
def general_conv(x,
                 num_filters=64,
                 filter_size=7,
                 stride=1,
                 stddev=0.02,
                 padding="VALID",
                 name="conv",
                 do_norm="instance",
                 do_relu=True,
                 relufactor=0):
  """"""
  with tf.variable_scope(name):
    x = layers().Conv2D(
        num_filters,
        filter_size,
        stride,
        padding,
        activation=None,
        kernel_initializer=tf.truncated_normal_initializer(stddev=stddev),
        bias_initializer=tf.constant_initializer(0.0))(x)
    if do_norm == "layer":
      x = layer_norm(x)
    elif do_norm == "instance":
      x = instance_norm(x)

    if do_relu:
      if relufactor == 0:
        x = tf.nn.relu(x, "relu")
      else:
        x = lrelu(x, leak=relufactor)

    return x
37439
function createComputations1000to1(n, sources) {
    for (var i = 0; i < n; i++) {
        createComputation1000(sources, i * 1000);
    }
}
856000
def sliced_gan():
  """"""
  hparams = common_hparams.basic_params1()
  hparams.optimizer = "adam"
  hparams.learning_rate_constant = 0.0002
  hparams.learning_rate_warmup_steps = 500
  hparams.learning_rate_schedule = "constant * linear_warmup"
  hparams.label_smoothing = 0.0
  hparams.batch_size = 128
  hparams.hidden_size = 128
  hparams.initializer = "uniform_unit_scaling"
  hparams.initializer_gain = 1.0
  hparams.weight_decay = 1e-6
  hparams.kernel_height = 4
  hparams.kernel_width = 4
  hparams.bottleneck_bits = 128
  hparams.add_hparam("discriminator_batchnorm", True)
  hparams.add_hparam("num_sliced_vecs", 4096)
  return hparams
1037130
def rnn(name, input, state, kernel, bias, new_state, number_of_gates = 2):
    ''' 
    '''

    nn = Build(name)
    nn.tanh(
        nn.mad(kernel=kernel, bias=bias,
            x=nn.concat(input, state)),
        out=new_state);
    return nn.layers;
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)