import tensorflow as tf

    means_norm_sq = tf.reduce_sum(tf.square(means), axis=-1, keep_dims=True)
    scalar_prod = tf.matmul(
        tf.transpose(x, perm=[1, 0, 2]), tf.transpose(means, perm=[0, 2, 1]))
    scalar_prod = tf.transpose(scalar_prod, perm=[1, 0, 2])
