import tensorflow as tf

        return tf.transpose(split_states(x, n), [0, 2, 1, 3])

def merge_heads(x):
    #[-1,head,n_ctx,emb]
    return merge_states(tf.transpose(x, [0, 2, 1, 3]))

def conv1d(x, scope, nf, rf, w_init=tf.random_normal_initializer(stddev=0.02), b_init=tf.constant_initializer(0), pad='VALID', train=False):
    with tf.variable_scope(scope):
        #x = [-1,n_ctx,512]
        nx = shape_list(x)[-1]
        #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)
