bb9b23bb210fbe363687d992e769bb1b23c21107,magenta/models/nsynth/wavenet/fastgen.py,,sample_categorical,#Any#,31

Before Change

def sample_categorical(pmf):
  cdf = np.cumsum(pmf)
  idx = np.random.rand()
  i = cdf.searchsorted(idx)
  return i


def load_nsynth(batch_size=1, sample_length=64000):

After Change

def sample_categorical(pmf):
  if pmf.ndim == 1:
    pmf = np.expand_dims(pmf, 0)
  batch_size = pmf.shape[0]
  cdf = np.cumsum(pmf, axis=1)
  rand_vals = np.random.rand(batch_size)
  idxs = np.zeros([batch_size, 1])
  for i in range(batch_size):
    idxs[i] = cdf[i].searchsorted(rand_vals[i])
  return idxs

Instances