Python tensorflow.floormod() Examples

def int_to_bit(x_int, num_bits, base=2):
  """Turn x_int representing numbers into a bitwise (lower-endian) tensor.

  Args:
    x_int: Tensor containing integer to be converted into base notation.
    num_bits: Number of bits in the representation.
    base: Base of the representation.

  Returns:
    Corresponding number expressed in base.
  """
  x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
  x_labels = []
  for i in range(num_bits):
    x_labels.append(
        tf.floormod(
            tf.floordiv(tf.to_int32(x_l),
                        tf.to_int32(base)**i), tf.to_int32(base)))
  res = tf.concat(x_labels, axis=-1)
  return tf.to_float(res) 

def _upsample_rois(scores, bboxes, keep_top_k):
    # upsample with replacement
    # filter out paddings
    bboxes = tf.boolean_mask(bboxes, scores > 0.)
    scores = tf.boolean_mask(scores, scores > 0.)

    scores, bboxes = tf.cond(tf.less(tf.shape(scores)[0], 1), lambda: (tf.constant([1.]), tf.constant([[0.2, 0.2, 0.8, 0.8]])), lambda: (scores, bboxes))
    #scores = tf.Print(scores,[scores])
    def upsampel_impl():
        num_bboxes = tf.shape(scores)[0]
        left_count = keep_top_k - num_bboxes

        select_indices = tf.random_shuffle(tf.range(num_bboxes))[:tf.floormod(left_count, num_bboxes)]
        #### zero
        select_indices = tf.concat([tf.tile(tf.range(num_bboxes), [tf.floor_div(left_count, num_bboxes) + 1]), select_indices], axis = 0)

        return [tf.gather(scores, select_indices), tf.gather(bboxes, select_indices)]
    return tf.cond(tf.shape(scores)[0] < keep_top_k, lambda : upsampel_impl(), lambda : [scores, bboxes]) 

def pyramidal_stack(outputs, sequence_length):
    shape = tf.shape(outputs)
    batch_size, max_time = shape[0], shape[1]
    num_units = outputs.get_shape().as_list()[-1]
    paddings = [[0, 0], [0, tf.floormod(max_time, 2)], [0, 0]]
    outputs = tf.pad(outputs, paddings)

    '''
    even_time = outputs[:, ::2, :]
    odd_time = outputs[:, 1::2, :]

    concat_outputs = tf.concat([even_time, odd_time], -1)
    '''

    concat_outputs = tf.reshape(outputs, (batch_size, -1, num_units * 2))

    return concat_outputs, tf.floordiv(sequence_length, 2) + tf.floormod(sequence_length, 2) 

def int_to_bit(x_int, num_bits, base=2):
  """Turn x_int representing numbers into a bitwise (lower-endian) tensor.

  Args:
    x_int: Tensor containing integer to be converted into base notation.
    num_bits: Number of bits in the representation.
    base: Base of the representation.

  Returns:
    Corresponding number expressed in base.
  """
  x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
  x_labels = [tf.floormod(
      tf.floordiv(tf.to_int32(x_l), tf.to_int32(base)**i), tf.to_int32(base))
              for i in range(num_bits)]
  res = tf.concat(x_labels, axis=-1)
  return tf.to_float(res) 

def int_to_bit(self, x_int, num_bits, base=2):

        """Turn x_int representing numbers into a bitwise (lower-endian)
        tensor.

        Args:
            x_int: Tensor containing integer to be converted into base
            notation.
            num_bits: Number of bits in the representation.
            base: Base of the representation.

        Returns:
            Corresponding number expressed in base.
        """
        x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
        x_labels = []
        for i in range(num_bits):
            x_labels.append(
                tf.floormod(
                    tf.floordiv(tf.to_int32(x_l),
                                tf.to_int32(base) ** i), tf.to_int32(base)))
        res = tf.concat(x_labels, axis=-1)
        return tf.to_float(res) 

def int_to_bit(x_int, num_bits, base=2):
  """Turn x_int representing numbers into a bitwise (lower-endian) tensor.

  Args:
    x_int: Tensor containing integer to be converted into base notation.
    num_bits: Number of bits in the representation.
    base: Base of the representation.

  Returns:
    Corresponding number expressed in base.
  """
  x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
  x_labels = []
  for i in range(num_bits):
    x_labels.append(
        tf.floormod(
            tf.floordiv(tf.to_int32(x_l),
                        tf.to_int32(base)**i), tf.to_int32(base)))
  res = tf.concat(x_labels, axis=-1)
  return tf.to_float(res) 

def int_to_bit(self, x_int, num_bits, base=2):
    """Turn x_int representing numbers into a bitwise (lower-endian) tensor.

    Args:
        x_int: Tensor containing integer to be converted into base
        notation.
        num_bits: Number of bits in the representation.
        base: Base of the representation.

    Returns:
        Corresponding number expressed in base.
    """
    x_l = tf.to_int32(tf.expand_dims(x_int, axis=-1))
    x_labels = []
    for i in range(num_bits):
      x_labels.append(
          tf.floormod(
              tf.floordiv(tf.to_int32(x_l),
                          tf.to_int32(base)**i), tf.to_int32(base)))
    res = tf.concat(x_labels, axis=-1)
    return tf.to_float(res) 

def test_FloorMod(self):
        t = tf.floormod(*self.random((4, 3), (4, 3)))
        self.check(t) 

def _test_forward_floormod(in_shape, if_shape, dtype):
    np_numer = np.random.uniform(1, 100, size=in_shape).astype(dtype)
    np_factor = np.random.uniform(1, 100, size=if_shape).astype(dtype)
    tf.reset_default_graph()
    with tf.Graph().as_default():
        numerator = tf.placeholder(dtype, in_shape, name="numer")
        factor = tf.placeholder(dtype, if_shape, name="factor")
        tf.floormod(numerator, factor, name='FloorMod')
        compare_tf_with_tvm([np_numer, np_factor], ['numer:0', 'factor:0'], 'FloorMod:0') 

def reset_gradient(global_step, accumulate_gradients, opt_reset):
    accu = tf.floormod(global_step, accumulate_gradients)
    if tf.equal(accu, 0):
        return opt_reset
    else:
        return 0.0 

def accumulate_gradient(global_step, accumulate_gradients, opt_compute, opt_apply):
    accu = tf.floormod(global_step, accumulate_gradients)
    if tf.equal(accu, 0):
        return opt_apply
    else:
        return opt_compute 

def get_keypoint(image, targets, predictions, heatmap_size, height, width, category, clip_at_zero=True, data_format='channels_last', name=None):
    predictions = tf.reshape(predictions, [1, -1, heatmap_size*heatmap_size])

    pred_max = tf.reduce_max(predictions, axis=-1)
    pred_indices = tf.argmax(predictions, axis=-1)
    pred_x, pred_y = tf.cast(tf.floormod(pred_indices, heatmap_size), tf.float32), tf.cast(tf.floordiv(pred_indices, heatmap_size), tf.float32)

    width, height = tf.cast(width, tf.float32), tf.cast(height, tf.float32)
    pred_x, pred_y = pred_x * width / tf.cast(heatmap_size, tf.float32), pred_y * height / tf.cast(heatmap_size, tf.float32)

    if clip_at_zero:
      pred_x, pred_y =  pred_x * tf.cast(pred_max>0, tf.float32), pred_y * tf.cast(pred_max>0, tf.float32)
      pred_x = pred_x * tf.cast(pred_max>0, tf.float32) + tf.cast(pred_max<=0, tf.float32) * (width / 2.)
      pred_y = pred_y * tf.cast(pred_max>0, tf.float32) + tf.cast(pred_max<=0, tf.float32) * (height / 2.)

    if config.PRED_DEBUG:
      pred_indices_ = tf.squeeze(pred_indices)
      image_ = tf.squeeze(image) * 255.
      pred_heatmap = tf.one_hot(pred_indices_, heatmap_size*heatmap_size, on_value=1., off_value=0., axis=-1, dtype=tf.float32)

      pred_heatmap = tf.reshape(pred_heatmap, [-1, heatmap_size, heatmap_size])
      if data_format == 'channels_first':
        image_ = tf.transpose(image_, perm=(1, 2, 0))
      save_image_op = tf.py_func(save_image_with_heatmap,
                                  [image_, height, width,
                                  heatmap_size,
                                  tf.reshape(pred_heatmap * 255., [-1, heatmap_size, heatmap_size]),
                                  tf.reshape(predictions, [-1, heatmap_size, heatmap_size]),
                                  config.left_right_group_map[category][0],
                                  config.left_right_group_map[category][1],
                                  config.left_right_group_map[category][2]],
                                  tf.int64, stateful=True)
      with tf.control_dependencies([save_image_op]):
        pred_x, pred_y = pred_x * 1., pred_y * 1.
    return pred_x, pred_y 

def load_cGAN_dataset(image_paths, semantic_map_paths, batch_size, test=False, augment=False, downsample=False,
            training_dataset='cityscapes'):
        """
        Load image dataset with semantic label maps for conditional GAN
        """ 

        def _parser(image_path, semantic_map_path):
            def _aspect_preserving_width_resize(image, width=512):
                # If training on ADE20k
                height_i = tf.shape(image)[0]
                new_height = height_i - tf.floormod(height_i, 16)
                    
                return tf.image.resize_image_with_crop_or_pad(image, new_height, width)

            def _image_decoder(path):
                im = tf.image.decode_png(tf.read_file(image_path), channels=3)
                im = tf.image.convert_image_dtype(im, dtype=tf.float32)
                return 2 * im - 1 # [0,1] -> [-1,1] (tanh range)


            image, semantic_map = _image_decoder(image_path), _image_decoder(semantic_map_path)
            
            print('Training on', training_dataset)
            if training_dataset is 'ADE20k':
                image = _aspect_preserving_width_resize(image)
                semantic_map = _aspect_preserving_width_resize(semantic_map)

            # im.set_shape([512,1024,3])  # downscaled cityscapes

            return image, semantic_map

        dataset = tf.data.Dataset.from_tensor_slices(image_paths, semantic_map_paths)
        dataset = dataset.map(_parser)
        dataset = dataset.shuffle(buffer_size=8)
        dataset = dataset.batch(batch_size)

        if test:
            dataset = dataset.repeat()

        return dataset 

def load_cGAN_dataset(image_paths, semantic_map_paths, batch_size, test=False, augment=False, downsample=False,
            training_dataset='cityscapes'):
        """
        Load image dataset with semantic label maps for conditional GAN
        """ 

        def _parser(image_path, semantic_map_path):
            def _aspect_preserving_width_resize(image, width=512):
                # If training on ADE20k
                height_i = tf.shape(image)[0]
                new_height = height_i - tf.floormod(height_i, 16)
                    
                return tf.image.resize_image_with_crop_or_pad(image, new_height, width)

            def _image_decoder(path):
                im = tf.image.decode_png(tf.read_file(image_path), channels=3)
                im = tf.image.convert_image_dtype(im, dtype=tf.float32)
                return 2 * im - 1 # [0,1] -> [-1,1] (tanh range)


            image, semantic_map = _image_decoder(image_path), _image_decoder(semantic_map_path)
            
            print('Training on', training_dataset)
            if training_dataset is 'ADE20k':
                image = _aspect_preserving_width_resize(image)
                semantic_map = _aspect_preserving_width_resize(semantic_map)

            # im.set_shape([512,1024,3])  # downscaled cityscapes

            return image, semantic_map

        dataset = tf.data.Dataset.from_tensor_slices(image_paths, semantic_map_paths)
        dataset = dataset.map(_parser)
        dataset = dataset.shuffle(buffer_size=8)
        dataset = dataset.batch(batch_size)

        if test:
            dataset = dataset.repeat()

        return dataset 

def one_hot_multiply(inputs, scale):
  """Performs (inputs * scale) % vocab_size in the one-hot space.

  Args:
    inputs: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
      Tensor.
    scale: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
      Tensor specifying how much to scale the corresponding one-hot vector in
      inputs. Soft values perform a "weighted scale": for example,
      scale=[0.2, 0.3, 0.5] performs a linear combination of
      0.2 * scaling by zero; 0.3 * scaling by one; and 0.5 * scaling by two.

  Returns:
    Tensor of same shape and dtype as inputs.
  """
  # TODO(trandustin): Implement with circular conv1d.
  inputs = tf.convert_to_tensor(inputs)
  scale = tf.cast(scale, inputs.dtype)
  batch_shape = inputs.shape[:-1].as_list()
  vocab_size = inputs.shape[-1].value
  # Form a [..., vocab_size, vocab_size] tensor. The ith row of the
  # batched vocab_size x vocab_size matrix represents scaling inputs by i.
  permutation_matrix = tf.floormod(
      tf.tile(tf.range(vocab_size)[:, tf.newaxis], [1, vocab_size]) *
      tf.range(vocab_size)[tf.newaxis], vocab_size)
  permutation_matrix = tf.one_hot(permutation_matrix, depth=vocab_size, axis=-1)
  # Scale the inputs according to the permutation matrix of all possible scales.
  scaled_inputs = tf.einsum('...v,avu->...au', inputs, permutation_matrix)
  scaled_inputs = tf.concat([tf.zeros(batch_shape + [1, vocab_size]),
                             scaled_inputs[..., 1:, :]], axis=-2)
  # Reduce rows of the scaled inputs by the scale values. This forms a
  # weighted linear combination of scaling by zero, scaling by one, and so on.
  outputs = tf.einsum('...v,...vu->...u', scale, scaled_inputs)
  return outputs 

def testMultiplicativeInverse(self):
    batch_size = 3
    vocab_size = 79
    length = 5
    inputs = np.random.randint(0, vocab_size - 1, size=(batch_size, length))
    one_hot_inputs = tf.one_hot(inputs, depth=vocab_size)

    one_hot_inv = reversible.multiplicative_inverse(one_hot_inputs, vocab_size)
    inv_inputs = tf.argmax(one_hot_inv, axis=-1)
    inputs_inv_inputs = tf.floormod(inputs * inv_inputs, vocab_size)
    inputs_inv_inputs_val = self.evaluate(inputs_inv_inputs)
    self.assertAllEqual(inputs_inv_inputs_val, np.ones((batch_size, length))) 

def ae_latent_softmax(latents_pred, latents_discrete, hparams):
  """Latent prediction and loss."""
  vocab_size = 2 ** hparams.z_size
  if hparams.num_decode_blocks < 2:
    latents_logits = tf.layers.dense(latents_pred, vocab_size,
                                     name="extra_logits")
    if hparams.logit_normalization:
      latents_logits *= tf.rsqrt(1e-8 +
                                 tf.reduce_mean(tf.square(latents_logits)))

    loss = None
    if latents_discrete is not None:
      if hparams.soft_em:
        # latents_discrete is actually one-hot of multinomial samples
        assert hparams.num_decode_blocks == 1
        loss = tf.nn.softmax_cross_entropy_with_logits_v2(
            labels=latents_discrete, logits=latents_logits)
      else:
        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
            labels=latents_discrete, logits=latents_logits)
    sample = multinomial_sample(
        latents_logits, vocab_size, hparams.sampling_temp)
    return sample, loss

  # Multi-block case.
  vocab_bits = int(math.log(vocab_size, 2))
  assert vocab_size == 2**vocab_bits
  assert vocab_bits % hparams.num_decode_blocks == 0
  block_vocab_size = 2**(vocab_bits // hparams.num_decode_blocks)
  latents_logits = [
      tf.layers.dense(
          latents_pred, block_vocab_size, name="extra_logits_%d" % i)
      for i in range(hparams.num_decode_blocks)
  ]
  loss = None
  if latents_discrete is not None:
    losses = []
    for i in range(hparams.num_decode_blocks):
      d = tf.floormod(tf.floordiv(latents_discrete,
                                  block_vocab_size**i), block_vocab_size)
      losses.append(tf.nn.sparse_softmax_cross_entropy_with_logits(
          labels=d, logits=latents_logits[i]))
    loss = sum(losses)
  samples = [multinomial_sample(l, block_vocab_size, hparams.sampling_temp)
             for l in latents_logits]
  sample = sum([s * block_vocab_size**i for i, s in enumerate(samples)])
  return sample, loss 

def ae_latent_softmax(latents_pred, latents_discrete, hparams):
  """Latent prediction and loss."""
  vocab_size = 2 ** hparams.z_size
  if hparams.num_decode_blocks < 2:
    latents_logits = tf.layers.dense(latents_pred, vocab_size,
                                     name="extra_logits")
    if hparams.logit_normalization:
      latents_logits *= tf.rsqrt(1e-8 +
                                 tf.reduce_mean(tf.square(latents_logits)))

    loss = None
    if latents_discrete is not None:
      if hparams.soft_em:
        # latents_discrete is actually one-hot of multinomial samples
        assert hparams.num_decode_blocks == 1
        loss = tf.nn.softmax_cross_entropy_with_logits_v2(
            labels=latents_discrete, logits=latents_logits)
      else:
        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
            labels=latents_discrete, logits=latents_logits)
    sample = multinomial_sample(
        latents_logits, vocab_size, hparams.sampling_temp)
    return sample, loss

  # Multi-block case.
  vocab_bits = int(math.log(vocab_size, 2))
  assert vocab_size == 2**vocab_bits
  assert vocab_bits % hparams.num_decode_blocks == 0
  block_vocab_size = 2**(vocab_bits // hparams.num_decode_blocks)
  latents_logits = [
      tf.layers.dense(
          latents_pred, block_vocab_size, name="extra_logits_%d" % i)
      for i in range(hparams.num_decode_blocks)
  ]
  loss = None
  if latents_discrete is not None:
    losses = []
    for i in range(hparams.num_decode_blocks):
      d = tf.floormod(tf.floordiv(latents_discrete,
                                  block_vocab_size**i), block_vocab_size)
      losses.append(tf.nn.sparse_softmax_cross_entropy_with_logits(
          labels=d, logits=latents_logits[i]))
    loss = sum(losses)
  samples = [multinomial_sample(l, block_vocab_size, hparams.sampling_temp)
             for l in latents_logits]
  sample = sum([s * block_vocab_size**i for i, s in enumerate(samples)])
  return sample, loss 

def make_split_document_labels(num_splits, dev_splits, test_splits):
    """
    Adapts tensorflow dataset to produce additional elements that indicate whether each datapoint is in train, dev,
    or test

    Particularly, splits the data into num_split folds, and censors the censored_split fold

    Parameters
    ----------
    num_splits integer in [0,100)
    dev_splits list of integers in [0,num_splits)
    test_splits list of integers in [0, num_splits)

    Returns
    -------
    fn: A function that can be used to map a dataset to censor some of the document labels.
    """

    def _tf_in1d(a, b):
        """
        Tensorflow equivalent of np.in1d(a,b)
        """
        a = tf.expand_dims(a, 0)
        b = tf.expand_dims(b, 1)
        return tf.reduce_any(tf.equal(a, b), 1)

    def _tf_scalar_a_in1d_b(a, b):
        """
        Tensorflow equivalent of np.in1d(a,b)
        """
        return tf.reduce_any(tf.equal(a, b))

    def fn(data):
        many_split = data['many_split']
        reduced_split = tf.floormod(many_split, num_splits)  # reduce the many splits to just num_splits

        in_dev = _tf_scalar_a_in1d_b(reduced_split, dev_splits)
        in_test = _tf_scalar_a_in1d_b(reduced_split, test_splits)
        in_train = tf.logical_not(tf.logical_or(in_dev, in_test))

        # in_dev = _tf_in1d(reduced_splits, dev_splits)
        # in_test = _tf_in1d(reduced_splits, test_splits)
        # in_train = tf.logical_not(tf.logical_or(in_dev, in_test))

        # code expects floats
        in_dev = tf.cast(in_dev, tf.float32)
        in_test = tf.cast(in_test, tf.float32)
        in_train = tf.cast(in_train, tf.float32)

        return {**data, 'in_dev': in_dev, 'in_test': in_test, 'in_train': in_train}

    return fn 

def make_split_document_labels(num_splits, dev_splits, test_splits):
    """
    Adapts tensorflow dataset to produce additional elements that indicate whether each datapoint is in train, dev,
    or test

    Particularly, splits the data into num_split folds, and censors the censored_split fold

    Parameters
    ----------
    num_splits integer in [0,100)
    dev_splits list of integers in [0,num_splits)
    test_splits list of integers in [0, num_splits)

    Returns
    -------
    fn: A function that can be used to map a dataset to censor some of the document labels.
    """
    def _tf_in1d(a,b):
        """
        Tensorflow equivalent of np.in1d(a,b)
        """
        a = tf.expand_dims(a, 0)
        b = tf.expand_dims(b, 1)
        return tf.reduce_any(tf.equal(a, b), 1)

    def _tf_scalar_a_in1d_b(a, b):
        """
        Tensorflow equivalent of np.in1d(a,b)
        """
        return tf.reduce_any(tf.equal(a, b))

    def fn(data):
        many_split = data['many_split']
        reduced_split = tf.floormod(many_split, num_splits)  # reduce the many splits to just num_splits

        in_dev = _tf_scalar_a_in1d_b(reduced_split, dev_splits)
        in_test = _tf_scalar_a_in1d_b(reduced_split, test_splits)
        in_train = tf.logical_not(tf.logical_or(in_dev, in_test))

        # in_dev = _tf_in1d(reduced_splits, dev_splits)
        # in_test = _tf_in1d(reduced_splits, test_splits)
        # in_train = tf.logical_not(tf.logical_or(in_dev, in_test))

        # code expects floats
        in_dev = tf.cast(in_dev, tf.float32)
        in_test = tf.cast(in_test, tf.float32)
        in_train = tf.cast(in_train, tf.float32)

        return {**data, 'in_dev': in_dev, 'in_test': in_test, 'in_train': in_train}

    return fn 

def normalized_error(targets, predictions, norm_value, visible, isvalid,
             bacth_size, num_keypoint, heatmap_size,
             train_image_size, clip_at_zero=True, name=None):

  with variable_scope.variable_scope(name, 'normalized_error', (targets, predictions, norm_value, visible, isvalid, bacth_size, num_keypoint, train_image_size, heatmap_size)):

    total = metric_variable([], dtypes.float32, name='total')
    count = metric_variable([], dtypes.float32, name='count')

    targets, predictions = tf.reshape(targets, [bacth_size, num_keypoint, -1]), tf.reshape(predictions, [bacth_size, num_keypoint, -1])

    pred_max = tf.reduce_max(predictions, axis=-1)
    pred_indices = tf.argmax(predictions, axis=-1)
    pred_x, pred_y = tf.floormod(pred_indices, heatmap_size) * train_image_size / heatmap_size, tf.floordiv(pred_indices, heatmap_size) * train_image_size / heatmap_size
    pred_x, pred_y = tf.cast(pred_x, tf.float32), tf.cast(pred_y, tf.float32)
    if clip_at_zero:
      pred_x, pred_y =  pred_x * tf.cast(pred_max>0, tf.float32), pred_y * tf.cast(pred_max>0, tf.float32)

    gt_indices = tf.argmax(targets, axis=-1)
    gt_x, gt_y = tf.floormod(gt_indices, heatmap_size) * train_image_size / heatmap_size, tf.floordiv(gt_indices, heatmap_size) * train_image_size / heatmap_size

    gt_x, gt_y = tf.cast(gt_x, tf.float32), tf.cast(gt_y, tf.float32)
    #print(gt_x,gt_y,pred_x,pred_y)
    #print(norm_value)
    #print(gt_x)
    #print(pred_x)
    dist = _safe_div(tf.pow(tf.pow(gt_x - pred_x, 2.) + tf.pow(gt_y - pred_y, 2.), .5), tf.expand_dims(norm_value, -1), 'norm_dist')

    #print(visible, isvalid)

    #dist = tf.cond(tf.equal(tf.shape(visible)[-1], tf.shape(isvalid)[-1]), lambda : tf.boolean_mask(dist, tf.logical_and(visible>0, isvalid>0)), lambda : dist)
    #print(dist)
    dist = tf.boolean_mask(dist, tf.logical_and(visible>0, isvalid>0))
    #dist = dist * tf.cast(tf.logical_and(visible>0, isvalid>0), tf.float32)

    update_total_op = state_ops.assign(total, math_ops.reduce_sum(dist))#assign_add #assign
    update_count_op = state_ops.assign(count, tf.cast(tf.shape(dist)[0], tf.float32))#assign_add #assign

    mean_t = _safe_div(total, count, 'value')
    update_op = _safe_div(update_total_op, update_count_op, 'update_op')

    return mean_t, update_op 

def ae_latent_softmax(latents_pred, latents_discrete, hparams):
  """Latent prediction and loss."""
  vocab_size = 2 ** hparams.z_size
  if hparams.num_decode_blocks < 2:
    latents_logits = tf.layers.dense(latents_pred, vocab_size,
                                     name="extra_logits")
    if hparams.logit_normalization:
      latents_logits *= tf.rsqrt(1e-8 +
                                 tf.reduce_mean(tf.square(latents_logits)))

    loss = None
    if latents_discrete is not None:
      if hparams.soft_em:
        # latents_discrete is actually one-hot of multinomial samples
        assert hparams.num_decode_blocks == 1
        loss = tf.nn.softmax_cross_entropy_with_logits_v2(
            labels=latents_discrete, logits=latents_logits)
      else:
        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
            labels=latents_discrete, logits=latents_logits)
    sample = multinomial_sample(
        latents_logits, vocab_size, hparams.sampling_temp)
    return sample, loss

  # Multi-block case.
  vocab_bits = int(math.log(vocab_size, 2))
  assert vocab_size == 2**vocab_bits
  assert vocab_bits % hparams.num_decode_blocks == 0
  block_vocab_size = 2**(vocab_bits // hparams.num_decode_blocks)
  latents_logits = [
      tf.layers.dense(
          latents_pred, block_vocab_size, name="extra_logits_%d" % i)
      for i in range(hparams.num_decode_blocks)
  ]
  loss = None
  if latents_discrete is not None:
    losses = []
    for i in range(hparams.num_decode_blocks):
      d = tf.floormod(tf.floordiv(latents_discrete,
                                  block_vocab_size**i), block_vocab_size)
      losses.append(tf.nn.sparse_softmax_cross_entropy_with_logits(
          labels=d, logits=latents_logits[i]))
    loss = sum(losses)
  samples = [multinomial_sample(l, block_vocab_size, hparams.sampling_temp)
             for l in latents_logits]
  sample = sum([s * block_vocab_size**i for i, s in enumerate(samples)])
  return sample, loss 

def get_opt(self, optimizer_fn, init_lr_dict, optimizer_type_dict, **kargs):

		self.init_lr_dict = init_lr_dict
		self.optimizer_type_dict = optimizer_type_dict
		self.optimizer_dict = {}

		self.alternate_order = kargs.get('alternate_order', list(self.init_lr_dict.keys()))
		print("==alternate order==", self.alternate_order)

		for key in self.alternate_order:
			init_lr = self.init_lr_dict[key]
			optimizer_type = self.optimizer_type_dict[key]
			if optimizer_type != 'radam' and key not in ['ebm_logz']:
				learning_rate = optimizer_fn.lr_decay_fn(init_lr, self.opt_config.num_train_steps, **kargs)
				learning_rate = optimizer_fn.warm_up(learning_rate, init_lr, **kargs)
				tf.logging.info("****** leanring rate warm up:%s ******", key)
			elif key == 'ebm_logz':
				tf.logging.info("****** ebm logz learning rate ******")
				if kargs.get('ebm_logz_update_circle', False):
					lr_ratio = tf.floormod(
										tf.train.get_or_create_global_step(),
										kargs.get('ebm_logz_update', 5),
										name="ebm_logz_update"
									)
					lr_ratio = tf.cast(tf.equal(tf.cast(lr_ratio, tf.int32), 0), tf.float32)
					tf.logging.info("****** learning_rate circle update ****** with %s circle", kargs.get('ebm_logz_update', 5))
				else:
					lr_ratio = 1.0
					tf.logging.info("****** normal learning_rate ******")
				if not kargs.get("use_tpu", False):
					tf.summary.scalar('{}_lr_ratio'.format(key), lr_ratio)
				learning_rate = init_lr * lr_ratio
			if not kargs.get("use_tpu", False):
				tf.summary.scalar('{}_learning_rate'.format(key), learning_rate)

			tf.logging.info("****** model:%s, optimizer: %s, learning_rate:%s", key, optimizer_type, str(init_lr))
			opt = optimizer_fn.optimizer_op(learning_rate, train_op=optimizer_type, **kargs)

			if kargs.get("use_tpu", False):
				tf.logging.info("***** Using tpu cross shard optimizer *****")
				opt = tf.contrib.tpu.CrossShardOptimizer(opt)
			self.optimizer_dict[key] = opt 

def get_opt(self, optimizer_fn, init_lr_dict, optimizer_type_dict, **kargs):

		self.init_lr_dict = init_lr_dict
		self.optimizer_type_dict = optimizer_type_dict
		self.optimizer_dict = {}

		self.alternate_order = kargs.get('alternate_order', list(self.init_lr_dict.keys()))
		print("==alternate order==", self.alternate_order)

		for key in self.alternate_order:
			init_lr = self.init_lr_dict[key]
			optimizer_type = self.optimizer_type_dict[key]
			if optimizer_type != 'radam' and key not in ['ebm_logz']:
				learning_rate = optimizer_fn.lr_decay_fn(init_lr, self.opt_config.num_train_steps, **kargs)
				learning_rate = optimizer_fn.warm_up(learning_rate, init_lr, **kargs)
				tf.logging.info("****** leanring rate warm up:%s ******", key)
			elif key == 'ebm_logz':
				tf.logging.info("****** ebm logz learning rate ******")
				if kargs.get('ebm_logz_update_circle', False):
					lr_ratio = tf.floormod(
										tf.train.get_or_create_global_step(),
										kargs.get('ebm_logz_update', 5),
										name="ebm_logz_update"
									)
					lr_ratio = tf.cast(tf.equal(tf.cast(lr_ratio, tf.int32), 0), tf.float32)
					tf.logging.info("****** learning_rate circle update ****** with %s circle", kargs.get('ebm_logz_update', 5))
				else:
					lr_ratio = 1.0
					tf.logging.info("****** normal learning_rate ******")
				if not kargs.get("use_tpu", False):
					tf.summary.scalar('{}_lr_ratio'.format(key), lr_ratio)
				learning_rate = init_lr * lr_ratio
			if not kargs.get("use_tpu", False):
				tf.summary.scalar('{}_learning_rate'.format(key), learning_rate)

			tf.logging.info("****** model:%s, optimizer: %s, learning_rate:%s", key, optimizer_type, str(init_lr))
			opt = optimizer_fn.optimizer_op(learning_rate, train_op=optimizer_type, **kargs)

			if kargs.get("use_tpu", False):
				tf.logging.info("***** Using tpu cross shard optimizer *****")
				opt = tf.contrib.tpu.CrossShardOptimizer(opt)
			self.optimizer_dict[key] = opt 

def sample_ensemble(config):
	'''Samples a random transformation according to the config.
	   Uses Latin Hypercube Sampling when batch size is greater than 1.'''
	
	N = config.batch_size

	# Offset randomization.
	offset_xy = tf.random_uniform([N, 2], minval=0, maxval=config.offset_max + 1, dtype=tf.int32)
			
	# Sample scale level.
	cumulative_sum = np.cumsum(config.scale_probabilities)
	u = cumulative_sum[-1] * tf.random_uniform([])
		
	scale_level = switch_case_cond(
		[(tf.less(u, x), (lambda j=i: tf.constant(j+1))) for i, x in enumerate(cumulative_sum[:-1])],
		lambda: tf.constant(len(cumulative_sum))
	)
	scale_level = tf.clip_by_value(scale_level, 1, config.num_scales)		
	
	# Scale randomization.
	scale_offset_xy = tf.random_uniform([2], minval=0, maxval=scale_level, dtype=tf.int32)
	
	# Sample flips.
	flips = tf.range((N + 3)//4*4, dtype=tf.int32)
	flips = tf.floormod(flips, 4)
	flips = tf.random_shuffle(flips)
	flips = flips[:N]
		
	# Sample transposing.
	swap_xy = tf.random_uniform([], minval=0, maxval=2, dtype=tf.int32)

	# Color multiplication.
	def sample_colors():
		color = tf.random_uniform([N], minval=0.0, maxval=1.0, dtype=config.dtype)
		color += tf.cast(tf.range(N), config.dtype)
		color /= tf.cast(N, config.dtype)
		return tf.random_shuffle(color)
	colors_r = tf.reshape(sample_colors(), [-1, 1, 1, 1])
	colors_g = tf.reshape(sample_colors(), [-1, 1, 1, 1])
	colors_b = tf.reshape(sample_colors(), [-1, 1, 1, 1])
	
	if config.color_multiplication_mode == 'color':
		color_factors = tf.concat([colors_r, colors_g, colors_b], axis=3)
	elif config.color_multiplication_mode == 'brightness':
		color_factors = tf.concat([colors_r, colors_r, colors_r], axis=3)
	else:
		raise Exception('Unknown color multiplication mode.')
	
	color_factors = 0.2 + 0.8 * color_factors
	
	# Sample permutations.
	permutations = np.asarray(list(itertools.permutations(range(3))), dtype=np.int32)
	repeat_count = (N + len(permutations) - 1) // len(permutations)
	permutations = tf.tile(tf.convert_to_tensor(permutations), tf.constant([repeat_count, 1]))
	permutations = tf.reshape(tf.random_shuffle(permutations)[:N, :], [-1])
			
	base_indices = 3 * tf.reshape(tf.tile(tf.reshape(tf.range(N), [-1, 1]), [1, 3]), [-1]) # [0, 0, 0, 3, 3, 3, 6, 6, 6, ...]
	permutations += base_indices
						
	return (offset_xy, flips, swap_xy, color_factors, permutations, scale_offset_xy, scale_level)