def test_ignore_logprob_basic():
x = Normal.dist()
(measurable_x_out, ) = get_measurable_outputs(x.owner.op, x.owner)
assert measurable_x_out is x.owner.outputs[1]
new_x = ignore_logprob(x)
assert new_x is not x
assert isinstance(new_x.owner.op, Normal)
assert type(new_x.owner.op).__name__ == "UnmeasurableNormalRV"
# Confirm that it does not have measurable output
assert get_measurable_outputs(new_x.owner.op, new_x.owner) is None
# Test that it will not clone a variable that is already unmeasurable
new_new_x = ignore_logprob(new_x)
assert new_new_x is new_x
def dist(
cls, mu=0.0, sigma=1.0, *, init=None, steps=None, size=None, **kwargs
) -> at.TensorVariable:
mu = at.as_tensor_variable(floatX(mu))
sigma = at.as_tensor_variable(floatX(sigma))
steps = get_steps(
steps=steps,
shape=kwargs.get("shape", None),
step_shape_offset=1,
)
if steps is None:
raise ValueError("Must specify steps or shape parameter")
steps = at.as_tensor_variable(intX(steps))
# If no scalar distribution is passed then initialize with a Normal of same mu and sigma
if init is None:
init = Normal.dist(0, 1)
else:
if not (
isinstance(init, at.TensorVariable)
and init.owner is not None
and isinstance(init.owner.op, RandomVariable)
and init.owner.op.ndim_supp == 0
):
raise TypeError("init must be a univariate distribution variable")
check_dist_not_registered(init)
# Ignores logprob of init var because that's accounted for in the logp method
init = ignore_logprob(init)
return super().dist([mu, sigma, init, steps], size=size, **kwargs)
def dist(
cls,
dist,
lower=None,
upper=None,
size=None,
shape=None,
**kwargs,
):
cls._argument_checks(dist, **kwargs)
lower, upper, initval = cls._set_values(lower,
upper,
size,
shape,
initval=None)
dist = ignore_logprob(dist)
if isinstance(dist.owner.op, Continuous):
res = _ContinuousBounded.dist(
[dist, lower, upper],
size=size,
shape=shape,
**kwargs,
)
res.tag.test_value = floatX(initval)
else:
res = _DiscreteBounded.dist(
[dist, lower, upper],
size=size,
shape=shape,
**kwargs,
)
res.tag.test_value = intX(initval)
return res
def dist(
cls,
rho,
sigma=None,
tau=None,
*,
init_dist=None,
steps=None,
constant=False,
ar_order=None,
**kwargs,
):
_, sigma = get_tau_sigma(tau=tau, sigma=sigma)
sigma = at.as_tensor_variable(floatX(sigma))
rhos = at.atleast_1d(at.as_tensor_variable(floatX(rho)))
if "init" in kwargs:
warnings.warn(
"init parameter is now called init_dist. Using init will raise an error in a future release.",
FutureWarning,
)
init_dist = kwargs.pop("init")
ar_order = cls._get_ar_order(rhos=rhos, constant=constant, ar_order=ar_order)
steps = get_steps(steps=steps, shape=kwargs.get("shape", None), step_shape_offset=ar_order)
if steps is None:
raise ValueError("Must specify steps or shape parameter")
steps = at.as_tensor_variable(intX(steps), ndim=0)
if init_dist is not None:
if not isinstance(init_dist, TensorVariable) or not isinstance(
init_dist.owner.op, RandomVariable
):
raise ValueError(
f"Init dist must be a distribution created via the `.dist()` API, "
f"got {type(init_dist)}"
)
check_dist_not_registered(init_dist)
if init_dist.owner.op.ndim_supp > 1:
raise ValueError(
"Init distribution must have a scalar or vector support dimension, ",
f"got ndim_supp={init_dist.owner.op.ndim_supp}.",
)
else:
warnings.warn(
"Initial distribution not specified, defaulting to "
"`Normal.dist(0, 100, shape=...)`. You can specify an init_dist "
"manually to suppress this warning.",
UserWarning,
)
init_dist = Normal.dist(0, 100, shape=(*sigma.shape, ar_order))
# Tell Aeppl to ignore init_dist, as it will be accounted for in the logp term
init_dist = ignore_logprob(init_dist)
return super().dist([rhos, sigma, init_dist, steps, ar_order, constant], **kwargs)
def test_ignore_logprob_model():
# logp that does not depend on input
def logp(value, x):
return value
with Model() as m:
x = Normal.dist()
y = DensityDist("y", x, logp=logp)
# Aeppl raises a KeyError when it finds an unexpected RV
with pytest.raises(KeyError):
joint_logp([y], {y: y.type()})
with Model() as m:
x = ignore_logprob(Normal.dist())
y = DensityDist("y", x, logp=logp)
assert joint_logp([y], {y: y.type()})
def __new__(
cls,
name,
dist,
lower=None,
upper=None,
size=None,
shape=None,
initval=None,
dims=None,
**kwargs,
):
cls._argument_checks(dist, **kwargs)
if dims is not None:
model = modelcontext(None)
if dims in model.coords:
dim_obj = np.asarray(model.coords[dims])
size = dim_obj.shape
else:
raise ValueError(
"Given dims do not exist in model coordinates.")
lower, upper, initval = cls._set_values(lower, upper, size, shape,
initval)
dist = ignore_logprob(dist)
if isinstance(dist.owner.op, Continuous):
res = _ContinuousBounded(
name,
[dist, lower, upper],
initval=floatX(initval),
size=size,
shape=shape,
**kwargs,
)
else:
res = _DiscreteBounded(
name,
[dist, lower, upper],
initval=intX(initval),
size=size,
shape=shape,
**kwargs,
)
return res
def dist(cls, mu=0.0, sigma=1.0, *, init_dist=None, steps=None, **kwargs) -> at.TensorVariable:
mu = at.as_tensor_variable(floatX(mu))
sigma = at.as_tensor_variable(floatX(sigma))
steps = get_steps(
steps=steps,
shape=kwargs.get("shape"),
step_shape_offset=1,
)
if steps is None:
raise ValueError("Must specify steps or shape parameter")
steps = at.as_tensor_variable(intX(steps))
if "init" in kwargs:
warnings.warn(
"init parameter is now called init_dist. Using init will raise an error in a future release.",
FutureWarning,
)
init_dist = kwargs.pop("init")
# If no scalar distribution is passed then initialize with a Normal of same mu and sigma
if init_dist is None:
warnings.warn(
"Initial distribution not specified, defaulting to `Normal.dist(0, 100)`."
"You can specify an init_dist manually to suppress this warning.",
UserWarning,
)
init_dist = Normal.dist(0, 100)
else:
if not (
isinstance(init_dist, at.TensorVariable)
and init_dist.owner is not None
and isinstance(init_dist.owner.op, RandomVariable)
and init_dist.owner.op.ndim_supp == 0
):
raise TypeError("init must be a univariate distribution variable")
check_dist_not_registered(init_dist)
# Ignores logprob of init var because that's accounted for in the logp method
init_dist = ignore_logprob(init_dist)
return super().dist([mu, sigma, init_dist, steps], **kwargs)
def rv_op(cls, weights, *components, size=None):
# Create new rng for the mix_indexes internal RV
mix_indexes_rng = aesara.shared(np.random.default_rng())
single_component = len(components) == 1
ndim_supp = components[0].owner.op.ndim_supp
if size is not None:
components = cls._resize_components(size, *components)
elif not single_component:
# We might need to broadcast components when size is not specified
shape = tuple(at.broadcast_shape(*components))
size = shape[:len(shape) - ndim_supp]
components = cls._resize_components(size, *components)
# Extract replication ndims from components and weights
ndim_batch = components[0].ndim - ndim_supp
if single_component:
# One dimension is taken by the mixture axis in the single component case
ndim_batch -= 1
# The weights may imply extra batch dimensions that go beyond what is already
# implied by the component dimensions (ndim_batch)
weights_ndim_batch = max(0, weights.ndim - ndim_batch - 1)
# If weights are large enough that they would broadcast the component distributions
# we try to resize them. This in necessary to avoid duplicated values in the
# random method and for equivalency with the logp method
if weights_ndim_batch:
new_size = at.concatenate([
weights.shape[:weights_ndim_batch],
components[0].shape[:ndim_batch],
])
components = cls._resize_components(new_size, *components)
# Extract support and batch ndims from components and weights
ndim_batch = components[0].ndim - ndim_supp
if single_component:
ndim_batch -= 1
weights_ndim_batch = max(0, weights.ndim - ndim_batch - 1)
assert weights_ndim_batch == 0
# Component RVs terms are accounted by the Mixture logprob, so they can be
# safely ignored by Aeppl
components = [ignore_logprob(component) for component in components]
# Create a OpFromGraph that encapsulates the random generating process
# Create dummy input variables with the same type as the ones provided
weights_ = weights.type()
components_ = [component.type() for component in components]
mix_indexes_rng_ = mix_indexes_rng.type()
mix_axis = -ndim_supp - 1
# Stack components across mixture axis
if single_component:
# If single component, we consider it as being already "stacked"
stacked_components_ = components_[0]
else:
stacked_components_ = at.stack(components_, axis=mix_axis)
# Broadcast weights to (*batched dimensions, stack dimension), ignoring support dimensions
weights_broadcast_shape_ = stacked_components_.shape[:ndim_batch + 1]
weights_broadcasted_ = at.broadcast_to(weights_,
weights_broadcast_shape_)
# Draw mixture indexes and append (stack + ndim_supp) broadcastable dimensions to the right
mix_indexes_ = at.random.categorical(weights_broadcasted_,
rng=mix_indexes_rng_)
mix_indexes_padded_ = at.shape_padright(mix_indexes_, ndim_supp + 1)
# Index components and squeeze mixture dimension
mix_out_ = at.take_along_axis(stacked_components_,
mix_indexes_padded_,
axis=mix_axis)
mix_out_ = at.squeeze(mix_out_, axis=mix_axis)
# Output mix_indexes rng update so that it can be updated in place
mix_indexes_rng_next_ = mix_indexes_.owner.outputs[0]
mix_op = MarginalMixtureRV(
inputs=[mix_indexes_rng_, weights_, *components_],
outputs=[mix_indexes_rng_next_, mix_out_],
)
# Create the actual MarginalMixture variable
mix_out = mix_op(mix_indexes_rng, weights, *components)
# Reference nodes to facilitate identification in other classmethods
mix_out.tag.weights = weights
mix_out.tag.components = components
mix_out.tag.choices_rng = mix_indexes_rng
return mix_out