code: def sample_prior(batch_size):
      cat, _ = get_distributions(DIST_PRIOR_PARAM[:NUM_CLASS], DIST_PRIOR_PARAM[NUM_CLASS:])
      sample_cat = tf.one_hot(cat.sample(batch_size), NUM_CLASS)
  
      
      sample_uni = tf.random_uniform([batch_size, NUM_UNIFORM], -1, 1)
      samples = tf.concat([sample_cat, sample_uni], axis=1)
      return samples
docs: OpenAI official code actually models the "uniform" latent code as
      a Gaussian distribution, but obtain the samples from a uniform distribution.