code: def direction(layer, vec, batch=None, cossim_pow=0):
    
    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)
docs: Visualize a direction