import tensorflow as tf

    if hparams.mode == tf.estimator.ModeKeys.PREDICT:
      layer_shape = common_layers.shape_list(layer)
      if hparams.full_latent_tower:
        rand = tf.random_uniform(layer_shape[:-1] + [hparams.bottleneck_bits])
      else:
        rand = tf.random_uniform(layer_shape[:-3] + [
