Visualize a direction (cossine similarity)

def direction_cossim(layer, vec, batch=None):
  """Visualize a direction (cossine similarity)"""
  def inner(T):
    act_mags = tf.sqrt(tf.reduce_sum(T(layer)**2, -1, keepdims=True))
    vec_mag = tf.sqrt(tf.reduce_sum(vec**2))
    mags = act_mags * vec_mag
    if batch is None:
      return tf.reduce_mean(T(layer) * vec.reshape([1, 1, 1, -1]) / mags)
    else:
      return tf.reduce_mean(T(layer)[batch] * vec.reshape([1, 1, -1]) / mags)
  return inner