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)

def sample_mog(prob, mean, var, rng):
    """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)
    """
    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