import tensorflow as tf

        #rf = 1,nx=emb,nf=3*emb
        w = tf.get_variable("w", [rf, nx, nf], initializer=w_init)
        b = tf.get_variable("b", [nf], initializer=b_init)
        if rf == 1: #faster 1x1 conv
            c = tf.reshape(tf.matmul(tf.reshape(x, [-1, nx]), tf.reshape(w, [-1, nf]))+b, shape_list(x)[:-1]+[nf])
        else: #was used to train LM
            c = tf.nn.conv1d(x, w, stride=1, padding=pad)+b
        return c

def attn(x, scope, n_state, n_head, train=False, scale=False):
    assert n_state%n_head==0
    with tf.variable_scope(scope):
        #c [-1,n_ctx,3*emb]
        c = conv1d(x, 'c_attn', n_state*3, 1, train=train)
        #q,k,v [-1,n_ctx,emb]
        q, k, v = tf.split(c, 3, 2)
        #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):
