def findMode_Mult_sort(Nvec, Wmat):
  ''' Find modes to multinomial with given parameters
        vectorized to solve many problems simultaneously

      Args
      ------
      Nvec : 1D array, size P
             contains non-negative integers
      Wmat : 2D array, size P x K
             each row is valid probability vector of length K
             Wmat[p] has non-negative entries, sums to one

      Returns
      -------
      Nmat : 2D sparse array, size P x K
          Nmat[p,:] = argmax_{n} \log p_{mult}( Nmat[p] | Nvec[p], Wmat[p] ) 
  '''
  # Unpack and parse input
  Nvec = np.asarray(Nvec, dtype=np.int64)
  if Nvec.ndim < 1:
    Nvec = Nvec[np.newaxis]
  Wmat = np.asarray(Wmat, dtype=np.float64)
  if Wmat.ndim < 2:
    Wmat = Wmat[np.newaxis, :]

  # Allocate and create initial guess, 
  #  which is guaranteed to have rows that sum to less than Nvec
  Nmat = np.zeros(Wmat.shape, dtype=np.int64)
  np.floor(Nvec[:,np.newaxis] * Wmat, out=Nmat)

  # Run Alg. 1 from paper
  activeErrors = Nvec - Nmat.sum(axis=1)
  activeRows = np.flatnonzero(activeErrors)

  if len(activeRows) > 0:
    sortIDs = np.argsort(-1 * activeErrors)
    np.take(Wmat, sortIDs, axis=0, out=Wmat)
    np.take(Nmat, sortIDs, axis=0, out=Nmat)    
    Nvec = Nvec[sortIDs]
    activeErrors = activeErrors[sortIDs]
    activeRows = np.flatnonzero(activeErrors)

    f = Nvec[:,np.newaxis] * Wmat - Nmat
    q = (1 - f)/ Wmat
    q = np.take(q, activeRows, axis=0)
    activeErrors = np.take(activeErrors, activeRows)    
    
  while len(activeRows) > 0:
    minIDs = np.argmin(q, axis=1)
    ids = np.ravel_multi_index([activeRows, minIDs], Nmat.shape)
    Nmat.ravel()[ids] += 1
    activeErrors -= 1

    stillActiveBound = np.searchsorted(-1*activeErrors, 0, 'left')

    activeErrors = activeErrors[:stillActiveBound]
    activeRows = activeRows[:stillActiveBound]
    ids = ids[:stillActiveBound]
    minIDs = minIDs[:stillActiveBound]
    q = q[:stillActiveBound]

    qids = np.ravel_multi_index([np.arange(len(minIDs)), minIDs], q.shape)
    q.ravel()[qids] += 1.0 / Wmat.ravel()[ids]
  
  # Double-check Nmat satisfies constraints          
  assert np.all(0.0001 < np.abs(Nmat.sum(axis=1) - Nvec))
  
  # Now just "unsort" the Nmat matrix
  unsortIDs = np.zeros(len(sortIDs), dtype=int)
  unsortIDs[sortIDs] = np.arange(len(sortIDs), dtype=int)
  np.take( Nmat, unsortIDs, axis=0, out=Nmat)
  return Nmat
dataset/ETHPy150Open daeilkim/refinery/refinery/bnpy/bnpy-dev/bnpy/util/NumericHardUtil.py/findMode_Mult_sort/swapped operands of `<`