def block()

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


def block(x, scope, *, past, mask, do_dropout, scale=False, hparams, seed):
    with tf.variable_scope(scope):
        attn_seed, mlp_seed = split_seed(seed, 2)

        nx = x.shape[-1].value
        a, present = attn(
            norm(x, 'ln_1'),
            'attn', nx, past=past, mask=mask, do_dropout=do_dropout, scale=scale, hparams=hparams, seed=attn_seed)
        x = x + a

        m = mlp(
            norm(x, 'ln_2'),
            'mlp', nx*4, do_dropout=do_dropout, hparams=hparams, seed=mlp_seed)
        h = x + m
        return h, present