def attention_2d()

in utils.py [0:0]

• 46 lines of code
• 2 McCabe index (conditional complexity)

def attention_2d(
        inp,
        q,
        q_b,
        k,
        k_b,
        v,
        v_b,
        reuse,
        scope,
        stop_at_grad=False,
        seperate=False,
        scale=False):
    inp_shape = tf.shape(inp)
    inp_compact = tf.reshape(
        inp,
        (inp_shape[0] *
         FLAGS.input_objects *
         inp_shape[1],
         inp.shape[3]))
    f_q = tf.matmul(inp_compact, q) + q_b
    f_k = tf.matmul(inp_compact, k) + k_b
    f_v = tf.nn.leaky_relu(tf.matmul(inp_compact, v) + v_b)

    f_q = tf.reshape(f_q,
                     (inp_shape[0],
                      inp_shape[1],
                         inp_shape[2],
                         tf.shape(f_q)[-1]))
    f_k = tf.reshape(f_k,
                     (inp_shape[0],
                      inp_shape[1],
                         inp_shape[2],
                         tf.shape(f_k)[-1]))
    f_v = tf.reshape(
        f_v,
        (inp_shape[0],
         inp_shape[1],
         inp_shape[2],
         inp_shape[3]))

    s = tf.matmul(f_k, f_q, transpose_b=True)
    c_num = (32**0.5)

    if scale:
        s = s / c_num

    beta = tf.nn.softmax(s, axis=-1)

    o = tf.reshape(tf.matmul(beta, f_v), inp_shape) + inp

    return o