def attn()

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

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

def attn(x, scope, n_state, *, past, mask, do_dropout, scale=False, hparams, seed):
    assert x.shape.ndims == 3  # Should be [batch, sequence, features]
    if past is not None:
        assert past.shape.ndims == 5  # Should be [batch, 2, heads, sequence, features], where 2 is [k, v]

    def split_heads(x):
        # From [batch, sequence, features] to [batch, heads, sequence, features]
        return tf.transpose(split_states(x, hparams.n_head), [0, 2, 1, 3])

    def merge_heads(x):
        # Reverse of split_heads
        return merge_states(tf.transpose(x, [0, 2, 1, 3]))

    def mask_attn_weights(w):
        # w has shape [batch, heads, dst_sequence, src_sequence], where information flows from src to dst.
        bs, _, nd, ns = utils.shape_list(w)
        b = attention_mask(nd, ns, dtype=w.dtype)
        b = tf.reshape(b, [1, 1, nd, ns])
        if mask is not None:
            b *= tf.reshape(tf.cast(mask, w.dtype), [bs, 1, 1, ns])
        w = w*b - tf.cast(1e10, w.dtype)*(1-b)
        return w

    def multihead_attn(q, k, v, *, seed):
        orig_dtype = v.dtype
        q, k, v = map(partial(tf.cast, dtype=tf.float32), (q, k, v))
        # q, k, v have shape [batch, heads, sequence, features]
        w = tf.matmul(q, k, transpose_b=True)

        if scale:
            n_state = v.shape[-1].value
            w = w * tf.rsqrt(tf.cast(n_state, w.dtype))

        w = mask_attn_weights(w)
        w = softmax(w)
        w = dropout(w, hparams.attn_pdrop,
                    do_dropout=do_dropout, name='attn_drop', stateless=True, seed=seed)
        a = tf.matmul(w, v)
        a = tf.cast(a, dtype=orig_dtype, name='a_cast')
        return a

    with tf.variable_scope(scope):
        attn_seed, resid_seed = split_seed(seed, 2)

        assert n_state % hparams.n_head == 0
        w_init_stdev = 1/np.sqrt(n_state)
        c = conv1x1(x, 'c_attn', n_state * 3, w_init_stdev=w_init_stdev)
        q, k, v = map(split_heads, tf.split(c, 3, axis=2))
        present = tf.stack([k, v], axis=1)
        if past is not None:
            pk, pv = tf.unstack(past, axis=1)
            k = tf.concat([pk, k], axis=-2)
            v = tf.concat([pv, v], axis=-2)
        a = multihead_attn(q, k, v, seed=attn_seed)
        a = merge_heads(a)
        w_init_stdev = 1/np.sqrt(n_state*hparams.n_layer)
        a = conv1x1(a, 'c_proj', n_state, w_init_stdev=w_init_stdev)
        a = dropout(a, hparams.resid_pdrop, do_dropout=do_dropout, stateless=True, seed=resid_seed, name='attn_resid_drop')
        return a, present