This initialization scheme follows Stan: we sample every latent
independently, uniformly from -2 to 2 in its unconstrained space,
and then transform into constrained space to construct an initial
state that can be passed to sample_chain or other MCMC drivers.
The argument signature is arranged to let the user pass either a
JointDistribution describing their model, if it's in that form, or
the essential information necessary for the sampling, namely a
bijector (from unconstrained to constrained space) and the desired
shape and dtype of each sample (specified in constrained space).
Args
model
A Distribution (typically a JointDistribution) giving the
model to be initialized. If supplied, it is queried for
its default event space bijector, its event shape, and its dtype.
If not supplied, those three elements must be supplied instead.
constraining_bijector
A (typically multipart) Bijector giving
the mapping from unconstrained to constrained space. If
supplied together with a model, acts as an override. A nested
structure of Bijectors is accepted, and interpreted as
applying in parallel to a corresponding structure of state parts
(see JointMap for details).
event_shapes
A structure of shapes giving the (unconstrained)
event space shape of the desired samples. Must be an acceptable
input to constraining_bijector.inverse_event_shape. If
supplied together with model, acts as an override.
event_shape_tensors
A structure of tensors giving the (unconstrained)
event space shape of the desired samples. Must be an acceptable
input to constraining_bijector.inverse_event_shape_tensor. If
supplied together with model, acts as an override. Required if any of
event_shapes are not fully-defined.
batch_shapes
A structure of shapes giving the batch shape of the desired
samples. If supplied together with model, acts as an override. If
unspecified, we assume scalar batch [].
batch_shape_tensors
A structure of tensors giving the batch shape of the
desired samples. If supplied together with model, acts as an override.
Required if any of batch_shapes are not fully-defined.
dtypes
A structure of dtypes giving the (unconstrained) dtypes of
the desired samples. Must be an acceptable input to
constraining_bijector.inverse_dtype. If supplied together
with model, acts as an override.
shard_axis_names
A structure of strs indicating the named axes by which
the distribution event is sharded. See
tfp.experimental.distribute.Sharded for more context.
Returns
init_dist
A Distribution representing the initialization
distribution, in constrained space. Samples from this
Distribution are valid initial states for a Markov chain
targeting the model.
Example
Initialize 100 chains from the unconstrained -2, 2 distribution
for a model expressed as a JointDistributionCoroutine:
[null,null,["Last updated 2023-11-21 UTC."],[],[],null,["# tfp.experimental.mcmc.init_near_unconstrained_zero\n\n\u003cbr /\u003e\n\n|------------------------------------------------------------------------------------------------------------------------------------------------------------|\n| [View source on GitHub](https://github.com/tensorflow/probability/blob/v0.23.0/tensorflow_probability/python/experimental/mcmc/initialization.py#L40-L209) |\n\nReturns an initialization Distribution for starting a Markov chain. \n\n tfp.experimental.mcmc.init_near_unconstrained_zero(\n model=None,\n constraining_bijector=None,\n event_shapes=None,\n event_shape_tensors=None,\n batch_shapes=None,\n batch_shape_tensors=None,\n dtypes=None,\n shard_axis_names=None\n )\n\n### Used in the notebooks\n\n| Used in the tutorials |\n|----------------------------------------------------------------------------------------------------------------------|\n| - [TFP Release Notes notebook (0.13.0)](https://www.tensorflow.org/probability/examples/TFP_Release_Notebook_0_13_0) |\n\nThis initialization scheme follows Stan: we sample every latent\nindependently, uniformly from -2 to 2 in its unconstrained space,\nand then transform into constrained space to construct an initial\nstate that can be passed to `sample_chain` or other MCMC drivers.\n\nThe argument signature is arranged to let the user pass either a\n`JointDistribution` describing their model, if it's in that form, or\nthe essential information necessary for the sampling, namely a\nbijector (from unconstrained to constrained space) and the desired\nshape and dtype of each sample (specified in constrained space).\n| **Note:** As currently implemented, this function has the limitation that the batch shape of the supplied model is ignored, but that could probably be generalized if needed.\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Args ---- ||\n|-------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n| `model` | A `Distribution` (typically a `JointDistribution`) giving the model to be initialized. If supplied, it is queried for its default event space bijector, its event shape, and its dtype. If not supplied, those three elements must be supplied instead. |\n| `constraining_bijector` | A (typically multipart) `Bijector` giving the mapping from unconstrained to constrained space. If supplied together with a `model`, acts as an override. A nested structure of `Bijector`s is accepted, and interpreted as applying in parallel to a corresponding structure of state parts (see `JointMap` for details). |\n| `event_shapes` | A structure of shapes giving the (unconstrained) event space shape of the desired samples. Must be an acceptable input to `constraining_bijector.inverse_event_shape`. If supplied together with `model`, acts as an override. |\n| `event_shape_tensors` | A structure of tensors giving the (unconstrained) event space shape of the desired samples. Must be an acceptable input to `constraining_bijector.inverse_event_shape_tensor`. If supplied together with `model`, acts as an override. Required if any of `event_shapes` are not fully-defined. |\n| `batch_shapes` | A structure of shapes giving the batch shape of the desired samples. If supplied together with `model`, acts as an override. If unspecified, we assume scalar batch `[]`. |\n| `batch_shape_tensors` | A structure of tensors giving the batch shape of the desired samples. If supplied together with `model`, acts as an override. Required if any of `batch_shapes` are not fully-defined. |\n| `dtypes` | A structure of dtypes giving the (unconstrained) dtypes of the desired samples. Must be an acceptable input to `constraining_bijector.inverse_dtype`. If supplied together with `model`, acts as an override. |\n| `shard_axis_names` | A structure of `str`s indicating the named axes by which the distribution event is sharded. See [`tfp.experimental.distribute.Sharded`](../../../tfp/experimental/distribute/Sharded) for more context. |\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Returns ------- ||\n|-------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n| `init_dist` | A `Distribution` representing the initialization distribution, in constrained space. Samples from this `Distribution` are valid initial states for a Markov chain targeting the model. |\n\n\u003cbr /\u003e\n\n#### Example\n\nInitialize 100 chains from the unconstrained -2, 2 distribution\nfor a model expressed as a `JointDistributionCoroutine`: \n\n @tfp.distributions.JointDistributionCoroutine\n def model():\n ...\n\n init_dist = tfp.experimental.mcmc.init_near_unconstrained_zero(model)\n states = tfp.mcmc.sample_chain(\n current_state=init_dist.sample(100, seed=[4, 8]),\n ...)"]]
View source on GitHub