def _call_execute_model(self, sample_shape, seed, value=None, sample_and_trace_fn=None): """Wraps the base `_call_execute_model` with vectorized_map.""" value_might_have_sample_dims = ( value is not None and _might_have_excess_ndims( # Double-flatten in case any components have structured events. flat_value=nest.flatten_up_to(self._single_sample_ndims, self._model_flatten(value), check_types=False), flat_core_ndims=tf.nest.flatten(self._single_sample_ndims))) sample_shape_may_be_nontrivial = ( distribution_util.shape_may_be_nontrivial(sample_shape)) if not self.use_vectorized_map or not (sample_shape_may_be_nontrivial or # pylint: disable=protected-access value_might_have_sample_dims): # No need to auto-vectorize. return joint_distribution_lib.JointDistribution._call_execute_model( # pylint: disable=protected-access self, sample_shape=sample_shape, seed=seed, value=value, sample_and_trace_fn=sample_and_trace_fn) # Set up for autovectorized sampling. To support the `value` arg, we need to # first understand which dims are from the model itself, then wrap # `_call_execute_model` to batch over all remaining dims. value_core_ndims = None if value is not None: value_core_ndims = tf.nest.map_structure( lambda v, nd: None if v is None else nd, value, self._model_unflatten(self._single_sample_ndims), check_types=False) vectorized_execute_model_helper = vectorization_util.make_rank_polymorphic( lambda v, seed: ( # pylint: disable=g-long-lambda joint_distribution_lib.JointDistribution._call_execute_model( # pylint: disable=protected-access self, sample_shape=(), seed=seed, value=v, sample_and_trace_fn=sample_and_trace_fn)), core_ndims=[value_core_ndims, None], validate_args=self.validate_args) # Redefine the polymorphic fn to hack around `make_rank_polymorphic` # not currently supporting keyword args. This is needed because the # `iid_sample` wrapper below expects to pass through a `seed` kwarg. vectorized_execute_model = ( lambda v, seed: vectorized_execute_model_helper(v, seed)) # pylint: disable=unnecessary-lambda if sample_shape_may_be_nontrivial: vectorized_execute_model = vectorization_util.iid_sample( vectorized_execute_model, sample_shape) return vectorized_execute_model(value, seed=seed)
def sample_distributions(self, sample_shape=(), seed=None, value=None, name='sample_distributions', **kwargs): """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"`. **kwargs: This is an alternative to passing a `value`, and achieves the same effect. Named arguments will be used to parameterize other dependent ("downstream") distribution-making functions. If a `value` argument is also provided, raises a ValueError. 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`. """ with self._name_and_control_scope(name): value = self._resolve_value(value=value, allow_partially_specified=True, **kwargs) might_have_batch_dims = ( distribution_util.shape_may_be_nontrivial(sample_shape) or value is not None) if self.use_vectorized_map and might_have_batch_dims: raise NotImplementedError( '`sample_distributions` with nontrivial ' 'sample shape is not yet supported ' 'for autovectorized JointDistributions.') ds, xs = self._call_flat_sample_distributions(sample_shape, seed=seed, value=value) if not might_have_batch_dims: # This is a single sample with no pinned values; this call will cache # the distributions if they are not already cached. self._get_single_sample_distributions(candidate_dists=ds) return self._model_unflatten(ds), self._model_unflatten(xs)
def state_space_model_likelihood(**param_vals): ssm = self.make_state_space_model( param_vals=param_vals, num_timesteps=num_timesteps, initial_step=initial_step, mask=mask, experimental_parallelize=experimental_parallelize) # Looping LGSSM methods are really expensive in eager mode; wrap them # to keep this from slowing things down in interactive use. ssm = tfe_util.JitPublicMethods(ssm, trace_only=True) if distribution_util.shape_may_be_nontrivial(trajectories_shape): return sample.Sample(ssm, sample_shape=trajectories_shape) return ssm
def _sample_n(self, sample_shape, seed, value=None): might_have_batch_dims = ( distribution_util.shape_may_be_nontrivial(sample_shape) or value is not None) if might_have_batch_dims: xs = self._call_execute_model( sample_shape, seed=seed, value=value, sample_and_trace_fn=trace_values_only) else: ds, xs = zip(*self._call_execute_model( sample_shape, seed=seed, value=value, sample_and_trace_fn=trace_distributions_and_values)) # This is a single sample with no pinned values; this call will cache # the distributions if they are not already cached. self._get_single_sample_distributions(candidate_dists=ds) return self._model_unflatten(xs)
def _affine_surrogate_posterior(event_shape, operators='diag', bijector=None, base_distribution=normal.Normal, dtype=tf.float32, batch_shape=(), validate_args=False, name=None): """Builds a joint variational posterior with a given `event_shape`. This function builds a surrogate posterior by applying a trainable transformation to a standard base distribution and constraining the samples with `bijector`. The surrogate posterior has event shape equal to the input `event_shape`. This function is a convenience wrapper around `build_affine_surrogate_posterior_from_base_distribution` that allows the user to pass in the desired posterior `event_shape` instead of pre-constructed base distributions (at the expense of full control over the base distribution types and parameterizations). Args: event_shape: (Nested) event shape of the posterior. operators: Either a string or a list/tuple containing `LinearOperator` subclasses, `LinearOperator` instances, or callables returning `LinearOperator` instances. Supported string values are "diag" (to create a mean-field surrogate posterior) and "tril" (to create a full-covariance surrogate posterior). A list/tuple may be passed to induce other posterior covariance structures. If the list is flat, a `tf.linalg.LinearOperatorBlockDiag` instance will be created and applied to the base distribution. Otherwise the list must be singly-nested and have a first element of length 1, second element of length 2, etc.; the elements of the outer list are interpreted as rows of a lower-triangular block structure, and a `tf.linalg.LinearOperatorBlockLowerTriangular` instance is created. For complete documentation and examples, see `tfp.experimental.vi.util.build_trainable_linear_operator_block`, which receives the `operators` arg if it is list-like. Default value: `"diag"`. bijector: `tfb.Bijector` instance, or nested structure of `tfb.Bijector` instances, that maps (nested) values in R^n to the support of the posterior. (This can be the `experimental_default_event_space_bijector` of the distribution over the prior latent variables.) Default value: `None` (i.e., the posterior is over R^n). base_distribution: A `tfd.Distribution` subclass parameterized by `loc` and `scale`. The base distribution of the transformed surrogate has `loc=0.` and `scale=1.`. Default value: `tfd.Normal`. dtype: The `dtype` of the surrogate posterior. Default value: `tf.float32`. batch_shape: Batch shape (Python tuple, list, or int) of the surrogate posterior, to enable parallel optimization from multiple initializations. Default value: `()`. validate_args: Python `bool`. Whether to validate input with asserts. This imposes a runtime cost. If `validate_args` is `False`, and the inputs are invalid, correct behavior is not guaranteed. Default value: `False`. name: Python `str` name prefixed to ops created by this function. Default value: `None` (i.e., 'build_affine_surrogate_posterior'). Yields: *parameters: sequence of `trainable_state_util.Parameter` namedtuples. These are intended to be consumed by `trainable_state_util.as_stateful_builder` and `trainable_state_util.as_stateless_builder` to define stateful and stateless variants respectively. #### Examples ```python tfd = tfp.distributions tfb = tfp.bijectors # Define a joint probabilistic model. Root = tfd.JointDistributionCoroutine.Root def model_fn(): concentration = yield Root(tfd.Exponential(1.)) rate = yield Root(tfd.Exponential(1.)) y = yield tfd.Sample( tfd.Gamma(concentration=concentration, rate=rate), sample_shape=4) model = tfd.JointDistributionCoroutine(model_fn) # Assume the `y` are observed, such that the posterior is a joint distribution # over `concentration` and `rate`. The posterior event shape is then equal to # the first two components of the model's event shape. posterior_event_shape = model.event_shape_tensor()[:-1] # Constrain the posterior values to be positive using the `Exp` bijector. bijector = [tfb.Exp(), tfb.Exp()] # Build a full-covariance surrogate posterior. surrogate_posterior = ( tfp.experimental.vi.build_affine_surrogate_posterior( event_shape=posterior_event_shape, operators='tril', bijector=bijector)) # For an example defining `'operators'` as a list to express an alternative # covariance structure, see # `build_affine_surrogate_posterior_from_base_distribution`. # Fit the model. y = [0.2, 0.5, 0.3, 0.7] target_model = model.experimental_pin(y=y) losses = tfp.vi.fit_surrogate_posterior( target_model.unnormalized_log_prob, surrogate_posterior, num_steps=100, optimizer=tf.optimizers.Adam(0.1), sample_size=10) ``` """ with tf.name_scope(name or 'build_affine_surrogate_posterior'): event_shape = nest.map_structure_up_to( _get_event_shape_shallow_structure(event_shape), lambda s: tf.convert_to_tensor(s, dtype=tf.int32), event_shape) if nest.is_nested(bijector): bijector = joint_map.JointMap(nest.map_structure( lambda b: identity.Identity() if b is None else b, bijector), validate_args=validate_args) if bijector is None: unconstrained_event_shape = event_shape else: unconstrained_event_shape = ( bijector.inverse_event_shape_tensor(event_shape)) standard_base_distribution = nest.map_structure( lambda s: base_distribution(loc=tf.zeros([], dtype=dtype), scale=1.), unconstrained_event_shape) standard_base_distribution = nest.map_structure( lambda d, s: ( # pylint: disable=g-long-lambda sample.Sample(d, sample_shape=s, validate_args=validate_args) if distribution_util.shape_may_be_nontrivial(s) else d), standard_base_distribution, unconstrained_event_shape) if distribution_util.shape_may_be_nontrivial(batch_shape): standard_base_distribution = nest.map_structure( lambda d: batch_broadcast.BatchBroadcast( # pylint: disable=g-long-lambda d, to_shape=batch_shape, validate_args=validate_args), standard_base_distribution) surrogate_posterior = yield from _affine_surrogate_posterior_from_base_distribution( standard_base_distribution, operators=operators, bijector=bijector, validate_args=validate_args) return surrogate_posterior