def direction(layer, vec, batch=None, cossim_pow=0):
  """Visualize a direction"""
  if batch is None:
    vec = vec[None, None, None]
    return lambda T: _dot_cossim(T(layer), vec)
  else:
    vec = vec[None, None]
    return lambda T: _dot_cossim(T(layer)[batch], vec)
Visualize a direction