def init_to_uniform(site=None, radius=2):
"""
Initialize to a random point in the area `(-radius, radius)` of unconstrained domain.
:param float radius: specifies the range to draw an initial point in the unconstrained domain.
"""
if site is None:
return partial(init_to_uniform, radius=radius)
if (site["type"] == "sample" and not site["is_observed"]
and not site["fn"].support.is_discrete):
if site["value"] is not None:
warnings.warn(
f"init_to_uniform() skipping initialization of site '{site['name']}'"
" which already stores a value.",
stacklevel=find_stack_level(),
)
return site["value"]
# XXX: we import here to avoid circular import
from numpyro.infer.util import helpful_support_errors
rng_key = site["kwargs"].get("rng_key")
sample_shape = site["kwargs"].get("sample_shape")
with helpful_support_errors(site):
transform = biject_to(site["fn"].support)
unconstrained_shape = transform.inverse_shape(site["fn"].shape())
unconstrained_samples = dist.Uniform(-radius, radius)(
rng_key=rng_key, sample_shape=sample_shape + unconstrained_shape)
return transform(unconstrained_samples)
def init_to_uniform(site=None, radius=2):
"""
Initialize to a random point in the area `(-radius, radius)` of unconstrained domain.
:param float radius: specifies the range to draw an initial point in the unconstrained domain.
"""
if site is None:
return partial(init_to_uniform, radius=radius)
if site['type'] == 'sample' and not site['is_observed'] and not site['fn'].is_discrete:
rng_key = site['kwargs'].get('rng_key')
sample_shape = site['kwargs'].get('sample_shape')
rng_key, subkey = random.split(rng_key)
# this is used to interpret the changes of event_shape in
# domain and codomain spaces
try:
prototype_value = site['fn'].sample(subkey, sample_shape=())
except NotImplementedError:
# XXX: this works for ImproperUniform prior,
# we can't use this logic for general priors
# because some distributions such as TransformedDistribution might
# have wrong event_shape.
prototype_value = jnp.full(site['fn'].shape(), jnp.nan)
transform = biject_to(site['fn'].support)
unconstrained_shape = jnp.shape(transform.inv(prototype_value))
unconstrained_samples = dist.Uniform(-radius, radius).sample(
rng_key, sample_shape=sample_shape + unconstrained_shape)
return transform(unconstrained_samples)
def __call__(self, name, fn, obs):
if name not in self.guide.prototype_trace:
return fn, obs
assert obs is None, "NeuTraReparam does not support observe statements"
log_density = 0.
if not self._x_unconstrained: # On first sample site.
# Sample a shared latent.
z_unconstrained = numpyro.sample(
"{}_shared_latent".format(self.guide.prefix),
self.guide.get_base_dist().mask(False))
# Differentiably transform.
x_unconstrained = self.transform(z_unconstrained)
# TODO: find a way to only compute those log_prob terms when needed
log_density = self.transform.log_abs_det_jacobian(
z_unconstrained, x_unconstrained)
self._x_unconstrained = self.guide._unpack_latent(x_unconstrained)
# Extract a single site's value from the shared latent.
unconstrained_value = self._x_unconstrained.pop(name)
transform = biject_to(fn.support)
value = transform(unconstrained_value)
logdet = transform.log_abs_det_jacobian(unconstrained_value, value)
logdet = sum_rightmost(
logdet,
jnp.ndim(logdet) - jnp.ndim(value) + len(fn.event_shape))
log_density = log_density + fn.log_prob(value) + logdet
numpyro.factor("_{}_log_prob".format(name), log_density)
return None, value
def gibbs_fn(rng_key, gibbs_sites, hmc_sites):
# convert to unconstrained values
z_hmc = {
k: biject_to(prototype_trace[k]["fn"].support).inv(v)
for k, v in hmc_sites.items()
if k in prototype_trace and prototype_trace[k]["type"] == "sample"
}
use_enum = len(set(support_sizes) - set(gibbs_sites)) > 0
wrapped_model = _wrap_model(model)
if use_enum:
from numpyro.contrib.funsor import config_enumerate, enum
wrapped_model = enum(config_enumerate(wrapped_model),
-max_plate_nesting - 1)
def potential_fn(z_discrete):
model_kwargs_ = model_kwargs.copy()
model_kwargs_["_gibbs_sites"] = z_discrete
return potential_energy(wrapped_model,
model_args,
model_kwargs_,
z_hmc,
enum=use_enum)
# get support_sizes of gibbs_sites
support_sizes_flat, _ = ravel_pytree(
{k: support_sizes[k]
for k in gibbs_sites})
num_discretes = support_sizes_flat.shape[0]
rng_key, rng_permute = random.split(rng_key)
idxs = random.permutation(rng_key, jnp.arange(num_discretes))
def body_fn(i, val):
idx = idxs[i]
support_size = support_sizes_flat[idx]
rng_key, z, pe = val
rng_key, z_new, pe_new, log_accept_ratio = proposal_fn(
rng_key,
z,
pe,
potential_fn=potential_fn,
idx=idx,
support_size=support_size)
rng_key, rng_accept = random.split(rng_key)
# u ~ Uniform(0, 1), u < accept_ratio => -log(u) > -log_accept_ratio
# and -log(u) ~ exponential(1)
z, pe = cond(
random.exponential(rng_accept) > -log_accept_ratio,
(z_new, pe_new), identity, (z, pe), identity)
return rng_key, z, pe
init_val = (rng_key, gibbs_sites, potential_fn(gibbs_sites))
_, gibbs_sites, _ = fori_loop(0, num_discretes, body_fn, init_val)
return gibbs_sites
def init_to_uniform(site=None, radius=2):
"""
Initialize to a random point in the area `(-radius, radius)` of unconstrained domain.
:param float radius: specifies the range to draw an initial point in the unconstrained domain.
"""
if site is None:
return partial(init_to_uniform, radius=radius)
if (site["type"] == "sample" and not site["is_observed"]
and not site["fn"].is_discrete):
rng_key = site["kwargs"].get("rng_key")
sample_shape = site["kwargs"].get("sample_shape")
rng_key, subkey = random.split(rng_key)
transform = biject_to(site["fn"].support)
unconstrained_shape = transform.inverse_shape(site["fn"].shape())
unconstrained_samples = dist.Uniform(-radius, radius)(
rng_key=rng_key, sample_shape=sample_shape + unconstrained_shape)
return transform(unconstrained_samples)
