docs: def safe_unit_norm(a):
      """
      Ensure that the vector or vectors have unit norm
      """
      if 1 == len(a.shape):
          n = np.linalg.norm(a)
          if n:
              return a / n
          return a
  
      norm = np.sum(np.abs(a) ** 2, axis=-1) ** (1. / 2)
  
      # Dividing by a norm of zero will cause a warning to be issued. Set those
      # values to another number. It doesn't matter what, since we'll be dividing
      # a vector of zeros by the number, and 0 / N always equals 0.
      norm[norm == 0] = -1e12
      return a / norm[:, np.newaxis]