def _sample_n(self, n, seed=None): seed = seed_stream.SeedStream(seed, "dirichlet_multinomial") n_draws = tf.cast(self.total_count, dtype=tf.int32) k = self.event_shape_tensor()[0] unnormalized_logits = tf.math.log( tf.random.gamma(shape=[n], alpha=self._broadcasted_concentration, dtype=self.dtype, seed=seed())) x = multinomial.draw_sample(1, k, unnormalized_logits, n_draws, self.dtype, seed()) final_shape = tf.concat([[n], self.batch_shape_tensor(), [k]], 0) return tf.reshape(x, final_shape)
def _sample_n(self, n, seed=None): # Need to create logits corresponding to [p, 1 - p]. # Note that for this distributions, logits corresponds to # inverse sigmoid(p) while in multivariate distributions, # such as multinomial this corresponds to log(p). # Because of this, when we construct the logits for the multinomial # sampler, we'll have to be careful. # log(p) = log(sigmoid(logits)) = logits - softplus(logits) # log(1 - p) = log(1 - sigmoid(logits)) = -softplus(logits) # Because softmax is invariant to a constnat shift in all inputs, # we can offset the logits by softplus(logits) so that we can use # [logits, 0.] as our input. logits = tf.stack([self.logits, tf.zeros_like(self.logits)], axis=-1) return multinomial.draw_sample(num_samples=n, num_classes=2, logits=logits, num_trials=tf.cast(self.total_count, dtype=tf.int32), dtype=self.dtype, seed=seed)[..., 0]
def _sample_n(self, n, seed=None): gamma_seed, multinomial_seed = samplers.split_seed( seed, salt='dirichlet_multinomial') concentration = tf.convert_to_tensor(self._concentration) total_count = tf.convert_to_tensor(self._total_count) n_draws = tf.cast(total_count, dtype=tf.int32) k = self._event_shape_tensor(concentration)[0] alpha = tf.math.multiply( tf.ones_like(total_count[..., tf.newaxis]), concentration, name='alpha') unnormalized_logits = gamma_lib.random_gamma( shape=[n], concentration=alpha, seed=gamma_seed, log_space=True) x = multinomial.draw_sample( 1, k, unnormalized_logits, n_draws, self.dtype, multinomial_seed) final_shape = ps.concat( [[n], self._batch_shape_tensor(concentration=concentration, total_count=total_count), [k]], 0) return tf.reshape(x, final_shape)
def _sample_n(self, n, seed=None): seed = seed_stream.SeedStream(seed, 'dirichlet_multinomial') concentration = tf.convert_to_tensor(self._concentration) total_count = tf.convert_to_tensor(self._total_count) n_draws = tf.cast(total_count, dtype=tf.int32) k = self._event_shape_tensor(concentration)[0] alpha = tf.math.multiply(tf.ones_like(total_count[..., tf.newaxis]), concentration, name='alpha') unnormalized_logits = tf.math.log( tf.random.gamma(shape=[n], alpha=alpha, dtype=self.dtype, seed=seed())) x = multinomial.draw_sample(1, k, unnormalized_logits, n_draws, self.dtype, seed()) final_shape = tf.concat( [[n], self._batch_shape_tensor(concentration, total_count), [k]], 0) return tf.reshape(x, final_shape)