code: def sample_categorical(probability_mass_function):
      """Sample from a categorical distribution.
  
      Args:
        probability_mass_function: Output of a softmax over categories.
          Array of shape [batch_size, number of categories]. Rows sum to 1.
  
      Returns:
        idxs: Array of size [batch_size, 1]. Integer of category sampled.
      """
      if probability_mass_function.ndim == 1:
          probability_mass_function = np.expand_dims(
              probability_mass_function, 0
          )
      batch_size = probability_mass_function.shape[0]
      cumulative_density_function = np.cumsum(
          probability_mass_function, axis=1
      )
      rand_vals = np.random.rand(batch_size)
      idxs = np.zeros([batch_size, 1])
      for i in range(batch_size):
          idxs[i] = cumulative_density_function[i].searchsorted(
              rand_vals[i]
          )
      return idxs
  
file_path: magenta/magenta/models/nsynth/wavenet/fastgen.py
imports: from magenta.models.nsynth import utils
  from magenta.models.nsynth.wavenet.h512_bo16 import Config
  from magenta.models.nsynth.wavenet.h512_bo16 import FastGenerationConfig
  import numpy
  from scipy.io import wavfile
  import tensorflow.compat.v1
label_desc: Use image node_id 1 for calling a global function with example usage: sample_categorical(probability_mass_function) and returns: idxs