code: def extend_dirichlet(p):
      
      if len(np.shape(p)) > 1:
          return np.hstack((p, np.atleast_2d(1. - np.sum(p))))
      else:
          return np.hstack((p, 1. - np.sum(p)))
docs: extend_dirichlet(p)
  
      Concatenates 1-sum(p) to the end of p and returns.