import tensorflow as tf


    if rank > 2:
        # reshape back to time dimension
        out = tf.reshape(out, shape=original_tensor_shape)

    return out


@layer
def dropout_layer(tensor, keep_prob=1.0, **opts):
    keep_prob = _global_keep_prob(keep_prob)
    out = tf.nn.dropout(tensor, keep_prob=keep_prob)
    return out


# TODO: should i normalize?
@layer
def word_dropout_layer(tensor, keep_prob=1.0, **opts):
    keep_prob = _global_keep_prob(keep_prob)

    rank = _rank(tensor)
    assert rank == 3, "Use embedding lookup layer"
