def dropout()

in lm_human_preferences/language/model.py [0:0]

• 16 lines of code
• 5 McCabe index (conditional complexity)

def dropout(x, pdrop, *, do_dropout, stateless=True, seed=None, name):
    """Like tf.nn.dropout but stateless.
    """
    if stateless:
        assert seed is not None
    def _dropout():
        with tf.name_scope(name):
            noise_shape = tf.shape(x)

            if stateless:
                r = tf.random.stateless_uniform(noise_shape, seed, dtype=x.dtype)
                # floor uniform [keep_prob, 1.0 + keep_prob)
                mask = tf.floor(1 - pdrop + r)
                return x * (mask * (1 / (1 - pdrop)))
            else:
                return tf.nn.dropout(x, rate=pdrop, noise_shape=noise_shape)
    if pdrop == 0 or not do_dropout:
        return x
    else:
        return _dropout()