code: def sample_mog(prob, mean, var, rng):
      
      gaussian_inds = sample_categorical(prob, rng).astype(numpy.int32)
      mean = mean[numpy.arange(mean.shape[0]), gaussian_inds, :]
      var = var[numpy.arange(mean.shape[0]), gaussian_inds, :]
      ret = sample_normal(mean=mean, var=var, rng=rng)
      return ret
docs: Sample from independent mixture of gaussian (MoG) distributions
  
      Each batch is an independent MoG distribution.
  
      Parameters
      ----------
      prob : numpy.ndarray
        mixture probability of each gaussian. Shape --> (batch_num, center_num)
      mean : numpy.ndarray
        mean of each gaussian. Shape --> (batch_num, center_num, sample_dim)
      var : numpy.ndarray
        variance of each gaussian. Shape --> (batch_num, center_num, sample_dim)
      rng : numpy.random.RandomState
  
      Returns
      -------
      ret : numpy.ndarray
        sampling result. Shape --> (batch_num, sample_dim)