import tensorflow as tf

        #q [-1,head,n_ctx,emb] v [-1,head,emb,n_ctx] v [-1,head,n_ctx,emb]
        q = split_heads(q, n_head)
        k = split_heads(k, n_head, k=True)
        v = split_heads(v, n_head)
        #a [-1,head,n_ctx,emb]
        a = _attn(q, k, v, train=train, scale=scale)
        #a [-1,n_ctx,head,emb]
        a = merge_heads(a)
        #a [-1,n_ctx,emb]
        a = conv1d(a, 'c_proj', n_state, 1, train=train)
        a = dropout(a, resid_pdrop, train)
        return a

def mlp(x, scope, n_state, train=False):
    with tf.variable_scope(scope):
        nx = shape_list(x)[-1]
        act = act_fns[afn]
        h = act(conv1d(x, 'c_fc', n_state, 1, train=train))
        h2 = conv1d(h, 'c_proj', nx, 1, train=train)
        h2 = dropout(h2, resid_pdrop, train)
        return h2

def block(x, scope, train=False, scale=False):
    with tf.variable_scope(scope):
        #nx = emb_size
        nx = shape_list(x)[-1]
        #a [-1,n_ctx,emb]
        a = attn(x, 'attn', nx, n_head, train=train, scale=scale)
