def _sample_n(self, n, seed): batch_shape = self.batch_shape_tensor() event_shape = self.event_shape_tensor() batch_ndims = tf.shape(batch_shape)[0] ndims = batch_ndims + 3 # sample_ndims=1, event_ndims=2 shape = tf.concat([[n], batch_shape, event_shape], 0) stream = seed_stream.SeedStream(seed, salt="Wishart") # Complexity: O(nbk**2) x = tf.random_normal(shape=shape, mean=0., stddev=1., dtype=self.dtype, seed=stream()) # Complexity: O(nbk) # This parametrization is equivalent to Chi2, i.e., # ChiSquared(k) == Gamma(alpha=k/2, beta=1/2) expanded_df = self.df * tf.ones( self.scale_operator.batch_shape_tensor(), dtype=self.df.dtype.base_dtype) g = tf.random_gamma(shape=[n], alpha=self._multi_gamma_sequence( 0.5 * expanded_df, self.dimension), beta=0.5, dtype=self.dtype, seed=stream()) # Complexity: O(nbk**2) x = tf.matrix_band_part(x, -1, 0) # Tri-lower. # Complexity: O(nbk) x = tf.matrix_set_diag(x, tf.sqrt(g)) # Make batch-op ready. # Complexity: O(nbk**2) perm = tf.concat([tf.range(1, ndims), [0]], 0) x = tf.transpose(x, perm) shape = tf.concat([batch_shape, [event_shape[0]], [-1]], 0) x = tf.reshape(x, shape) # Complexity: O(nbM) where M is the complexity of the operator solving a # vector system. For LinearOperatorLowerTriangular, each matmul is O(k^3) so # this step has complexity O(nbk^3). x = self.scale_operator.matmul(x) # Undo make batch-op ready. # Complexity: O(nbk**2) shape = tf.concat([batch_shape, event_shape, [n]], 0) x = tf.reshape(x, shape) perm = tf.concat([[ndims - 1], tf.range(0, ndims - 1)], 0) x = tf.transpose(x, perm) if not self.input_output_cholesky: # Complexity: O(nbk**3) x = tf.matmul(x, x, adjoint_b=True) return x
def _sample_n(self, n, seed=None): low = tf.convert_to_tensor(self.low) high = tf.convert_to_tensor(self.high) peak = tf.convert_to_tensor(self.peak) stream = seed_stream.SeedStream(seed, salt='triangular') shape = tf.concat( [[n], self._batch_shape_tensor(low=low, high=high, peak=peak)], axis=0) samples = tf.random.uniform(shape=shape, dtype=self.dtype, seed=stream()) # We use Inverse CDF sampling here. Because the CDF is a quadratic function, # we must use sqrts here. interval_length = high - low return tf.where( # Note the CDF on the left side of the peak is # (x - low) ** 2 / ((high - low) * (peak - low)). # If we plug in peak for x, we get that the CDF at the peak # is (peak - low) / (high - low). Because of this we decide # which part of the piecewise CDF we should use based on the cdf samples # we drew. samples < (peak - low) / interval_length, # Inverse of (x - low) ** 2 / ((high - low) * (peak - low)). low + tf.sqrt(samples * interval_length * (peak - low)), # Inverse of 1 - (high - x) ** 2 / ((high - low) * (high - peak)) high - tf.sqrt((1. - samples) * interval_length * (high - peak)))
def _sample_3d(self, n, seed=None): """Specialized inversion sampler for 3D.""" seed = seed_stream.SeedStream(seed, salt='von_mises_fisher_3d') u_shape = tf.concat([[n], self._batch_shape_tensor()], axis=0) z = tf.random.uniform(u_shape, seed=seed(), dtype=self.dtype) # TODO(bjp): Higher-order odd dim analytic CDFs are available in [1], could # be bisected for bounded sampling runtime (i.e. not rejection sampling). # [1]: Inversion sampler via: https://ieeexplore.ieee.org/document/7347705/ # The inversion is: u = 1 + log(z + (1-z)*exp(-2*kappa)) / kappa # We must protect against both kappa and z being zero. safe_conc = tf.where(self.concentration > 0, self.concentration, tf.ones_like(self.concentration)) safe_z = tf.where(z > 0, z, tf.ones_like(z)) safe_u = 1 + tf.reduce_logsumexp( [tf.math.log(safe_z), tf.math.log1p(-safe_z) - 2 * safe_conc], axis=0) / safe_conc # Limit of the above expression as kappa->0 is 2*z-1 u = tf.where(self.concentration > tf.zeros_like(safe_u), safe_u, 2 * z - 1) # Limit of the expression as z->0 is -1. u = tf.where(tf.equal(z, 0), -tf.ones_like(u), u) if not self._allow_nan_stats: u = tf.debugging.check_numerics(u, 'u in _sample_3d') return u[..., tf.newaxis]
def _sample_n(self, n, seed=None): # See https://en.wikipedia.org/wiki/Inverse_Gaussian_distribution or # https://www.jstor.org/stable/2683801 concentration = tf.convert_to_tensor(self.concentration) loc = tf.convert_to_tensor(self.loc) seed = seed_stream.SeedStream(seed, 'inverse_gaussian') shape = tf.concat( [[n], self._batch_shape_tensor(loc=loc, concentration=concentration)], axis=0) sampled_chi2 = (tf.random.normal(shape, mean=0., stddev=1., seed=seed(), dtype=self.dtype))**2. sampled_uniform = tf.random.uniform(shape, minval=0., maxval=1., seed=seed(), dtype=self.dtype) sampled = (loc + loc**2. * sampled_chi2 / (2. * concentration) - loc / (2. * concentration) * (4. * loc * concentration * sampled_chi2 + (loc * sampled_chi2)**2)**0.5) return tf.where(sampled_uniform <= loc / (loc + sampled), sampled, loc**2 / sampled)
def _sample_n(self, n, seed=None): seed = seed_stream.SeedStream(seed, salt="Autoregressive")() samples = self.distribution0.sample(n, seed=seed) for _ in range(self._num_steps): # pylint: disable=not-callable samples = self.distribution_fn(samples).sample(seed=seed) return samples
def _flat_sample_distributions(self, sample_shape=(), seed=None, value=None): """Executes `model`, creating both samples and distributions.""" ds = [] values_out = [] seed = seed_stream.SeedStream('JointDistributionCoroutine', seed) gen = self._model() index = 0 d = next(gen) try: while True: actual_distribution = d.distribution if isinstance(d, self.Root) else d ds.append(actual_distribution) if (value is not None and len(value) > index and value[index] is not None): seed() next_value = value[index] else: next_value = actual_distribution.sample( sample_shape=sample_shape if isinstance(d, self.Root) else (), seed=seed()) values_out.append(next_value) index += 1 d = gen.send(next_value) except StopIteration: pass return ds, values_out
def _sample_n(self, n, seed=None): # This implementation is equivalent to the one in tf.contrib.distributions.Wishart batch_shape = self.batch_shape_tensor() event_shape = self.event_shape_tensor() stream = seed_stream.SeedStream(seed=seed, salt="Wishart") shape = tf.concat([[n], batch_shape, event_shape], 0) # Sample a normal full matrix x = tf.random_normal(shape=shape, dtype=self.dtype, seed=stream()) # Sample the diagonal g = tf.random_gamma(shape=[n], alpha=self._multi_gamma_sequence( 0.5 * self.df, self.p), beta=0.5, dtype=self.dtype, seed=stream()) # Discard the upper triangular part x = tf.matrix_band_part(x, -1, 0) # Set the diagonal x = tf.matrix_set_diag(x, tf.sqrt(g)) # Scale with the Scale matrix, equivalent to matmul(sqrt(diag_scale), x) x *= tf.sqrt(tf.exp(self.log_diag_scale[tf.newaxis, :, :, tf.newaxis])) return x
def _flat_sample_distributions(self, sample_shape=(), seed=None, value=None): # This function additionally depends on: # self._dist_fn_wrapped # self._dist_fn_args # self._always_use_specified_sample_shape seed = seed_stream.SeedStream('JointDistributionSequential', seed) ds = [] xs = [None] * len(self._dist_fn_wrapped) if value is None else list( value) if len(xs) != len(self._dist_fn_wrapped): raise ValueError('Number of `xs`s must match number of ' 'distributions.') for i, (dist_fn, args) in enumerate( zip(self._dist_fn_wrapped, self._dist_fn_args)): ds.append(dist_fn(*xs[:i])) # Chain rule of probability. if xs[i] is None: # TODO(b/129364796): We should ignore args prefixed with `_`; this # would mean we more often identify when to use `sample_shape=()` # rather than `sample_shape=sample_shape`. xs[i] = ds[-1].sample( () if args and not self._always_use_specified_sample_shape else sample_shape, seed=seed()) else: xs[i] = tf.convert_to_tensor(value=xs[i], dtype_hint=ds[-1].dtype) seed( ) # Ensure reproducibility even when xs are (partially) set. # Note: we could also resolve distributions up to the first non-`None` in # `self._flatten(value)`, however we omit this feature for simplicity, # speed, and because it has not yet been requested. return ds, xs
def _sample_distributions(self, sample_shape=(), seed=None, value=None): seed = seed_stream.SeedStream('JointDistributionSequential', seed) ds = [] if value is None: xs = [None] * len(self._wrapped_dist_fn) else: xs = list(value) xs.extend([None] * (len(self._wrapped_dist_fn) - len(xs))) if len(xs) != len(self.distribution_fn): raise ValueError('Number of `xs`s must match number of ' 'distributions.') for i, (dist_fn, args) in enumerate( zip(self._wrapped_dist_fn, self._dist_fn_args)): ds.append(dist_fn(*xs[:i])) # Chain rule of probability. if xs[i] is None: # TODO(b/129364796): We should ignore args prefixed with `_`; this # would mean we more often identify when to use `sample_shape=()` # rather than `sample_shape=sample_shape`. xs[i] = ds[-1].sample( () if args and not self._always_use_specified_sample_shape else sample_shape, seed=seed()) else: xs[i] = tf.convert_to_tensor(value=xs[i], dtype_hint=ds[-1].dtype) seed( ) # Ensure reproducibility even when xs are (partially) set. self._most_recently_built_distributions = ds return tuple(ds), tuple(xs)
def _sample_n(self, n, seed=None): with tf.control_dependencies(self._assertions): seed = seed_stream.SeedStream(seed=seed, salt='BlockwiseDistribution') xs = [ tf.cast(d.sample(n, seed=seed()), self.dtype) for d in self.distributions ] return tf.concat(xs, axis=-1)
def _sample_n(self, n, seed): with tf.compat.v1.control_dependencies(self._runtime_assertions): seed = seed_stream.SeedStream(seed, salt="ZeroInflated") mask = self.inflated_distribution.sample(n, seed()) samples = self.count_distribution.sample(n, seed()) mask, samples = _broadcast_rate(mask, samples) # mask = 1 => new_sample = 0 # mask = 0 => new_sample = sample return samples * tf.cast(1 - mask, samples.dtype)
def _sample_n_sparse(self, n, kw, seed=None): assert n == 1 batch_shape = self.batch_shape_tensor() event_shape = self.event_shape iw = int(np.sqrt(event_shape[0].value)) # Image width nb = kw**2 # Number of basis nb_half = nb // 2 + 1 nch = 1 # Number of channels in the image stream = seed_stream.SeedStream(seed=seed, salt="Wishart") shape = tf.concat([batch_shape, [iw, iw, nb_half - 1]], 0) # Random sample for the off diagonal values as a dense tensor x_right = tf.random_normal(shape=shape, dtype=self.dtype, seed=stream()) # The upper triangular values needed to get a square kernel per pixel x_left = tf.zeros(shape) # Random sample for the diagonal of the matrix x_diag = tf.random_gamma(shape=[n], alpha=self._multi_gamma_sequence( 0.5 * self.df, self.p), beta=0.5, dtype=self.dtype, seed=stream()) # Concatenate the diagonal and off-diagonal elements x_diag = tf.reshape(x_diag, (-1, iw, iw, nch)) x = tf.concat([tf.sqrt(x_diag), x_right], axis=3) # Scale the sampled matrix using the distribution Scale matrix diag_scale = tf.exp(self.log_diag_scale) diag_scale = tf.reshape(diag_scale, (-1, iw, iw, nch)) x *= tf.sqrt(diag_scale) # Square root is equivalent to Cholesky # Concatenate with the zeros x = tf.concat([x_left, x], axis=3) # Create identity basis so that the sampled matrix is only defined by x, if this were not the case # we would have to do some optimization to find the basis and weights that reconstruct x identity_basis = tf.eye(num_rows=nb) identity_basis = tf.reshape(identity_basis, (nb, kw, kw, nch, nch)) sample_shape = tf.concat([batch_shape, [iw, iw, nch]], axis=0) x_sparse = cov_rep.PrecisionConvCholFilters( weights_precision=x, filters_precision=identity_basis, sample_shape=sample_shape) return x_sparse
def _sample_n(self, n, seed=None): seed = seed_stream.SeedStream(seed, "gamma_gamma") rate = tf.random_gamma(shape=[n], alpha=self.mixing_concentration, beta=self.mixing_rate, dtype=self.dtype, seed=seed()) return tf.random_gamma(shape=[], alpha=self.concentration, beta=rate, dtype=self.dtype, seed=seed())
def _sample_n(self, n, seed=None): seed = seed_stream.SeedStream(seed, "beta") concentration1 = tf.convert_to_tensor(self.concentration1) concentration0 = tf.convert_to_tensor(self.concentration0) shape = self._batch_shape_tensor(concentration1, concentration0) expanded_concentration1 = tf.broadcast_to(concentration1, shape) expanded_concentration0 = tf.broadcast_to(concentration0, shape) gamma1_sample = tf.random.gamma( shape=[n], alpha=expanded_concentration1, dtype=self.dtype, seed=seed()) gamma2_sample = tf.random.gamma( shape=[n], alpha=expanded_concentration0, dtype=self.dtype, seed=seed()) beta_sample = gamma1_sample / (gamma1_sample + gamma2_sample) return beta_sample
def _sample_n(self, n, seed=None): # Here we use the fact that if: # lam ~ Gamma(concentration=total_count, rate=(1-probs)/probs) # then X ~ Poisson(lam) is Negative Binomially distributed. stream = seed_stream.SeedStream(seed, salt="NegativeBinomial") rate = tf.random.gamma( shape=[n], alpha=self.total_count, beta=tf.exp(-self.logits), dtype=self.dtype, seed=stream()) return tf.random.poisson( lam=rate, shape=[], dtype=self.dtype, seed=stream())
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): seed = seed_stream.SeedStream(seed, "gamma_gamma") rate = tf.random.gamma( shape=[n], # Be sure to draw enough rates for the fully-broadcasted gamma-gamma. alpha=self.mixing_concentration + tf.zeros_like(self.concentration), beta=self.mixing_rate, dtype=self.dtype, seed=seed()) return tf.random.gamma(shape=[], alpha=self.concentration, beta=rate, dtype=self.dtype, seed=seed())
def sample_distributions(self, sample_shape=(), seed=None, value=None, name='sample_distributions'): """Generate samples and the (random) distributions. Note that a call to `sample()` without arguments will generate a single sample. Args: sample_shape: 0D or 1D `int32` `Tensor`. Shape of the generated samples. seed: Python integer seed for generating random numbers. value: `list` of `Tensor`s in `distribution_fn` order to use to parameterize other ("downstream") distribution makers. Default value: `None` (i.e., draw a sample from each distribution). name: name prepended to ops created by this function. Default value: `"sample_distributions"`. Returns: distributions: a `tuple` of `Distribution` instances for each of `distribution_fn`. samples: a `tuple` of `Tensor`s with prepended dimensions `sample_shape` for each of `distribution_fn`. """ seed = seed_stream.SeedStream('JointDistributionSequential', seed) with self._name_scope(name, values=[sample_shape, seed, value]): ds = [] if value is None: xs = [None] * len(self._wrapped_distribution_fn) else: xs = list(value) xs.extend([None]*(len(self._wrapped_distribution_fn) - len(xs))) if len(xs) != len(self.distribution_fn): raise ValueError('Number of `xs`s must match number of ' 'distributions.') for i, (dist_fn, args) in enumerate(zip(self._wrapped_distribution_fn, self._args)): ds.append(dist_fn(*xs[:i])) # Chain rule of probability. if xs[i] is None: # TODO(b/129364796): We should ignore args prefixed with `_`; this # would mean we more often identify when to use `sample_shape=()` # rather than `sample_shape=sample_shape`. xs[i] = ds[-1].sample( () if args and not self._always_use_specified_sample_shape else sample_shape, seed=seed()) else: xs[i] = tf.convert_to_tensor(value=xs[i], dtype_hint=ds[-1].dtype) seed() # Ensure reproducibility even when xs are (partially) set. self._most_recently_built_distributions = ds return tuple(ds), tuple(xs)
def _sample_n(self, n, seed=None): stream = seed_stream.SeedStream(seed=seed, salt="Wishart") # Sample a normal full matrix x = self.normal_dist.sample(sample_shape=n, seed=stream()) # Sample the log diagonal log_g = self.sqrt_gamma_dist.sample(sample_shape=n, seed=stream()) # Discard the upper triangular part x = tf.matrix_band_part(x, -1, 0) # Set the diagonal x = tf.matrix_set_diag(x, log_g) return x
def _sample_n(self, n, seed=None): seed = seed_stream.SeedStream(seed, "beta") expanded_concentration1 = tf.ones_like( self.total_concentration, dtype=self.dtype) * self.concentration1 expanded_concentration0 = tf.ones_like( self.total_concentration, dtype=self.dtype) * self.concentration0 gamma1_sample = tf.random.gamma(shape=[n], alpha=expanded_concentration1, dtype=self.dtype, seed=seed()) gamma2_sample = tf.random.gamma(shape=[n], alpha=expanded_concentration0, dtype=self.dtype, seed=seed()) beta_sample = gamma1_sample / (gamma1_sample + gamma2_sample) return beta_sample
def _sample_n(self, n, seed=None): # Like with the univariate Student's t, sampling can be implemented as a # ratio of samples from a multivariate gaussian with the appropriate # covariance matrix and a sample from the chi-squared distribution. seed = seed_stream.SeedStream(seed, salt="multivariate t") loc = tf.broadcast_to(self.loc, self._sample_shape()) mvn = mvn_linear_operator.MultivariateNormalLinearOperator( loc=tf.zeros_like(loc), scale=self.scale) normal_samp = mvn.sample(n, seed=seed()) df = tf.broadcast_to(self.df, self.batch_shape_tensor()) chi2 = chi2_lib.Chi2(df=df) chi2_samp = chi2.sample(n, seed=seed()) return (self._loc + normal_samp * tf.math.rsqrt(chi2_samp / self._df)[..., tf.newaxis])
def _sample_n(self, n, seed): with tf.control_dependencies(self._runtime_assertions): seed = seed_stream.SeedStream(seed, salt="MixtureSameFamily") x = self.components_distribution.sample(n, seed=seed()) # [n, B, k, E] # TODO(jvdillon): Consider using tf.gather (by way of index unrolling). npdt = x.dtype.as_numpy_dtype mask = tf.one_hot( indices=self.mixture_distribution.sample(n, seed=seed()), # [n, B] depth=self._num_components, # == k on_value=np.ones([], dtype=npdt), off_value=np.zeros([], dtype=npdt)) # [n, B, k] mask = distribution_utils.pad_mixture_dimensions( mask, self, self.mixture_distribution, self._event_ndims) # [n, B, k, [1]*e] x = tf.reduce_sum(x * mask, axis=-1 - self._event_ndims) # [n, B, E] if self._reparameterize: x = self._reparameterize_sample(x) return x
def _sample_n(self, n, seed=None): # The sampling method comes from the fact that if: # X ~ Normal(0, 1) # Z ~ Chi2(df) # Y = X / sqrt(Z / df) # then: # Y ~ StudentT(df). seed = seed_stream.SeedStream(seed, "student_t") shape = tf.concat([[n], self.batch_shape_tensor()], 0) normal_sample = tf.random.normal(shape, dtype=self.dtype, seed=seed()) df = self.df * tf.ones(self.batch_shape_tensor(), dtype=self.dtype) gamma_sample = tf.random.gamma([n], 0.5 * df, beta=0.5, dtype=self.dtype, seed=seed()) samples = normal_sample * tf.math.rsqrt(gamma_sample / df) return samples * self.scale + self.loc # Abs(scale) not wanted.
def _sample_n(self, n, seed=None): seed = seed_stream.SeedStream(seed, "normal_gamma") shape = tf.concat([[n], self.batch_shape_tensor()], 0) precision = tf.random_gamma( shape=shape, alpha=self.concentration, beta=self.rate, dtype=self.dtype, seed=seed()) scale = tf.sqrt(1 / (self._lambda * precision)) mean = tf.random_normal( shape=shape, mean=0., stddev=1., dtype=self.loc.dtype, seed=seed()) mean = mean * scale + self.loc return tf.concat((tf.expand_dims(mean, axis=-1), tf.expand_dims(precision, axis=-1)), axis=-1)
def _sample_n(self, n, seed=None): # Get ids as a [n, batch_size]-shaped matrix, unless batch_shape=[] then get # ids as a [n]-shaped vector. distributions = self.poisson_and_mixture_distributions() dist, mixture_dist = distributions batch_size = tensorshape_util.num_elements(self.batch_shape) if batch_size is None: batch_size = tf.reduce_prod( self._batch_shape_tensor(distributions=distributions)) # We need to 'sample extra' from the mixture distribution if it doesn't # already specify a probs vector for each batch coordinate. # We only support this kind of reduced broadcasting, i.e., there is exactly # one probs vector for all batch dims or one for each. stream = seed_stream.SeedStream( seed, salt='PoissonLogNormalQuadratureCompound') ids = mixture_dist.sample(sample_shape=concat_vectors( [n], distribution_util.pick_vector(mixture_dist.is_scalar_batch(), [batch_size], np.int32([]))), seed=stream()) # We need to flatten batch dims in case mixture_dist has its own # batch dims. ids = tf.reshape(ids, shape=concat_vectors([n], distribution_util.pick_vector( self.is_scalar_batch(), np.int32([]), np.int32([-1])))) # Stride `quadrature_size` for `batch_size` number of times. offset = tf.range(start=0, limit=batch_size * self._quadrature_size, delta=self._quadrature_size, dtype=ids.dtype) ids = ids + offset rate = tf.gather(tf.reshape(dist.rate, shape=[-1]), ids) rate = tf.reshape( rate, shape=concat_vectors( [n], self._batch_shape_tensor(distributions=distributions))) return tf.random.poisson(lam=rate, shape=[], dtype=self.dtype, seed=seed)
def _sample_n(self, n, seed=None): # See # https://en.wikipedia.org/wiki/Exponentially_modified_Gaussian_distribution shape = tf.concat([[n], self.batch_shape_tensor()], 0) stream = seed_stream.SeedStream(seed, salt="exponentially_modified_normal") sampled_normal = tf.random_normal(shape=shape, mean=0., stddev=1., dtype=self.loc.dtype, seed=stream()) sampled_uniform = tf.random_uniform( shape=shape, minval=np.finfo(self.dtype.as_numpy_dtype).tiny, maxval=1., dtype=self.dtype, seed=stream()) return (sampled_normal * self.scale + self.loc - tf.log(sampled_uniform) / self.rate)
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.reshape( tf.log(tf.random_gamma( shape=[n], alpha=self.concentration, dtype=self.dtype, seed=seed())), shape=[-1, k]) draws = tf.multinomial( logits=unnormalized_logits, num_samples=n_draws, seed=seed()) x = tf.reduce_sum(tf.one_hot(draws, depth=k), -2) final_shape = tf.concat([[n], self.batch_shape_tensor(), [k]], 0) x = tf.reshape(x, final_shape) return tf.cast(x, self.dtype)
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)
def test_seed_stream(salt='Salt of the Earth', hardcoded_seed=None): """Returns a command-line-controllable SeedStream PRNG for unit tests. When seeding unit-test PRNGs, we want: - The seed to be fixed to an arbitrary value most of the time, so the test doesn't flake even if its failure probability is noticeable. - To switch to different seeds per run when using --runs_per_test to measure the test's failure probability. - To set the seed to a specific value when reproducing a low-probability event (e.g., debugging a crash that only some seeds trigger). To those ends, this function returns a `SeedStream` seeded with `test_seed` (which see). The latter respects the command line flags `--fixed_seed=<seed>` and `--vary-seed` (Boolean, default False). `--vary_seed` uses system entropy to produce unpredictable seeds. `--fixed_seed` takes precedence over `--vary_seed` when both are present. Note that TensorFlow graph mode operations tend to read seed state from two sources: a "graph-level seed" and an "op-level seed". tf.test.TestCase will set the former to a fixed value per test, but in general it may be necessary to explicitly set both to ensure reproducibility. Args: salt: Optional string wherewith to salt the returned SeedStream. Setting this guarantees independent random numbers across tests. hardcoded_seed: Optional Python value. The seed to use if both the `--vary_seed` and `--fixed_seed` flags are unset. This should usually be unnecessary, since a test should pass with any seed. Returns: strm: A SeedStream instance seeded with 17, unless otherwise specified by arguments or command line flags. """ return seed_stream.SeedStream(salt, test_seed(hardcoded_seed))
def _flat_sample_distributions(self, sample_shape=(), seed=None, value=None): """Executes `model`, creating both samples and distributions.""" ds = [] values_out = [] seed = seed_stream.SeedStream('JointDistributionCoroutine', seed) gen = self._model() index = 0 d = next(gen) if not isinstance(d, self.Root): raise ValueError('First distribution yielded by coroutine must ' 'be wrapped in `Root`.') try: while True: actual_distribution = d.distribution if isinstance(d, self.Root) else d ds.append(actual_distribution) if (value is not None and len(value) > index and value[index] is not None): seed() next_value = value[index] else: next_value = actual_distribution.sample( sample_shape=sample_shape if isinstance(d, self.Root) else (), seed=seed()) if self._validate_args: with tf.control_dependencies( self._assert_compatible_shape( index, sample_shape, next_value)): values_out.append(tf.identity(next_value)) else: values_out.append(next_value) index += 1 d = gen.send(next_value) except StopIteration: pass return ds, values_out