Mean and attention to reduce spatial dimensions.

def mean_with_attention(x, name, num_heads=4):
  """Mean and attention to reduce spatial dimensions."""
  with tf.variable_scope(name):
    shape = shape_list(x)
    m = tf.reduce_mean(x, [1, 2])
    a = layers().Dense(num_heads, name="mean_attn")(x)
    s = tf.reshape(a, [shape[0], -1, num_heads])
    s = tf.nn.softmax(s, axis=1)
    s = tf.reshape(s, shape[:-1] + [1, num_heads])
    am = tf.reduce_mean(tf.expand_dims(x, axis=-1) * s, [1, 2])
    l = tf.concat([am, tf.expand_dims(m, axis=-1)], axis=-1)
    return layers().Dense(2 * shape[-1], name="mean_attn_final")(
        tf.reshape(l, [shape[0], (num_heads+1) * shape[-1]]))