code: def dist(x1, x2):
      
      x1p2 = np.sum(np.square(x1), 1)
      x2p2 = np.sum(np.square(x2), 1)
      return x1p2.reshape((-1, 1)) + x2p2.reshape((1, -1)) - 2 * np.dot(x1, x2.T)
docs: Compute squared euclidean distance between samples (autograd)