dynode.infer.sample.sample_distributions

dynode.infer.sample.sample_distributions(obj: Any, rng_key: Array | None = None, _prefix: str = '')

Recurisvely scans data structures and samples numpyro.Distribution objects.

Parameters

obj: Any

object to be sampled or searched for distributions.

rng_keyArray

optional rng_key to use if this function is called outside of an mcmc or inference context where one is automatically provided.

_prefixstr

prefix to prepend to all site names, used to build up site names.

Note

Sampled distributions receive site names according to a set of rules - Distributions within lists are appended with _i identifing the index, N times for N dimensional arrays - Dictionaries and Pydnatic models are recursively searched with sites names prepending the key that parameter belongs to.

Returns

Any | jax.Array

obj with all instances of numpyro.Distributions sampled, if obj itself is a distribution, the sample will be returned.