Model.sample()

Contents

Model.sample()#

Model.sample(shape, seed, posterior_samples=None, fixed=(), newdata=None, dists=None, chunk_size=64)[source]#

Draws samples from the model.

Parameters:

  • shape (Sequence[int]) – Sample shape.

  • seed (Array) – The seed is split and distributed to the seed nodes of the model. Must be a jax RNG key array that satisfies jnp.issubdtype(key.dtype, jax.dtypes.prng_key). See jax.random and https://docs.jax.dev/en/latest/jep/9263-typed-keys.html for more details.

  • posterior_samples (Position (dict[str, Any]) | None, default: None) – Position of samples at which to evaluate predictions. All values are assumed to have two leading dimensions corresponding to (nchains, niteration). Values are converted with their model-specific converters before sampling.

  • fixed (Sequence[str], default: ()) – The names of the nodes or variables to be excluded from the simulation. By default, no nodes or variables are skipped.

  • newdata (Position (dict[str, Any]) | None, default: None) – Position of new data at which to produce samples. The keys should correspond to variable or node names in the model whose values should be set to the given values before sampling. Values are converted with their model-specific converters. If None (default), the current variable values are used.

  • dists (dict[str, Dist] | None, default: None) – Can be used to provide a dictionary of variable names and Dist instances to use in sampling. If None (default), samples are drawn for each variable using their Var.dist_node.

  • chunk_size (int | None, default: 64) – Maximum number of flattened requested-draw and posterior-sample combinations to evaluate in parallel. Defaults to 64. Pass None to evaluate all combinations in parallel. A smaller value reduces the peak memory required for sample-dependent intermediate values, at the potential cost of lower accelerator utilization. It does not reduce the memory required to store the returned samples.

Notes

When compiling this function with jax.jit, the arguments shape, fixed, dists, and chunk_size must be static.

Return type:

Position (dict[str, Any])

Returns:

A dictionary of variable and node names and their sampled values. Includes only sampled variables.

Contents