code: def sample_categorical(prob, rng):
      
      ret = numpy.empty(prob.shape[0], dtype=numpy.float32)
      for ind in range(prob.shape[0]):
          ret[ind] = numpy.searchsorted(numpy.cumsum(prob[ind]), rng.rand()).clip(min=0.0,
                                                                                  max=prob.shape[
                                                                                          1] - 0.5)
      return ret
docs: Sample from independent categorical distributions
  
      Each batch is an independent categorical distribution.
  
      Parameters
      ----------
      prob : numpy.ndarray
        Probability of the categorical distribution. Shape --> (batch_num, category_num)
      rng : numpy.random.RandomState
  
      Returns
      -------
      ret : numpy.ndarray
        Sampling result. Shape --> (batch_num,)