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, w, comp_dists, **kwargs):
if not isinstance(comp_dists, (tuple, list)):
# comp_dists is a single component
comp_dists = [comp_dists]
elif len(comp_dists) == 1:
warnings.warn(
"Single component will be treated as a mixture across the last size dimension.\n"
"To disable this warning do not wrap the single component inside a list or tuple",
UserWarning,
)
# Check that components are not associated with a registered variable in the model
components_ndim_supp = set()
for dist in comp_dists:
# TODO: Allow these to not be a RandomVariable as long as we can call `ndim_supp` on them
# and resize them
if not isinstance(dist, TensorVariable) or not isinstance(
dist.owner.op, RandomVariable):
raise ValueError(
f"Component dist must be a distribution created via the `.dist()` API, got {type(dist)}"
)
check_dist_not_registered(dist)
components_ndim_supp.add(dist.owner.op.ndim_supp)
if len(components_ndim_supp) > 1:
raise ValueError(
f"Mixture components must all have the same support dimensionality, got {components_ndim_supp}"
)
w = at.as_tensor_variable(w)
return super().dist([w, *comp_dists], **kwargs)
def _argument_checks(cls, dist, **kwargs):
if "observed" in kwargs:
raise ValueError(
"Observed Bound distributions are not supported. "
"If you want to model truncated data "
"you can use a pm.Potential in combination "
"with the cumulative probability function."
)
if not isinstance(dist, TensorVariable):
raise ValueError(
"Passing a distribution class to `Bound` is no longer supported.\n"
"Please pass the output of a distribution instantiated via the "
"`.dist()` API such as:\n"
'`pm.Bound("bound", pm.Normal.dist(0, 1), lower=0)`'
)
check_dist_not_registered(dist)
if dist.owner.op.ndim_supp != 0:
raise NotImplementedError("Bounding of MultiVariate RVs is not yet supported.")
if not isinstance(dist.owner.op, (Discrete, Continuous)):
raise ValueError(
f"`distribution` {dist} must be a Discrete or Continuous" " distribution subclass"
)
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 dist(cls, dist, lower, upper, **kwargs):
if not isinstance(dist, TensorVariable) or not isinstance(dist.owner.op, RandomVariable):
raise ValueError(
f"Censoring dist must be a distribution created via the `.dist()` API, got {type(dist)}"
)
if dist.owner.op.ndim_supp > 0:
raise NotImplementedError(
"Censoring of multivariate distributions has not been implemented yet"
)
check_dist_not_registered(dist)
return super().dist([dist, lower, upper], **kwargs)
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)