import tensorflow as tf

_phase_train = _phase.assign(True)
_phase_infer = _phase.assign(False)


# TODO: move to ops
def _rank(x):
    return len(x.get_shape())


def _apply_dropout_mask(tensor_shape, keep_prob=1.0, normalize=True):
    random_tensor = keep_prob + tf.random_uniform(tensor_shape, dtype=tf.float32)
    binary_mask = tf.floor(random_tensor)
    if normalize:
        binary_mask = tf.reciprocal(keep_prob) * binary_mask
    return binary_mask


def _global_keep_prob(keep_prob):
    keep_prob = tf.convert_to_tensor(keep_prob, dtype=tf.float32)
    keep_prob = tf.cond(_phase, lambda: keep_prob, lambda: keep_prob * 0.0 + 1.0)
    return keep_prob


def layer(func):

    class Layer(object):
