def rv_op(cls, dist, lower=None, upper=None, size=None, rngs=None):
lower = at.constant(
-np.inf) if lower is None else at.as_tensor_variable(lower)
upper = at.constant(
np.inf) if upper is None else at.as_tensor_variable(upper)
# When size is not specified, dist may have to be broadcasted according to lower/upper
dist_shape = size if size is not None else at.broadcast_shape(
dist, lower, upper)
dist = change_rv_size(dist, dist_shape)
# Censoring is achieved by clipping the base distribution between lower and upper
rv_out = at.clip(dist, lower, upper)
# Reference nodes to facilitate identification in other classmethods, without
# worring about possible dimshuffles
rv_out.tag.dist = dist
rv_out.tag.lower = lower
rv_out.tag.upper = upper
if rngs is not None:
rv_out = cls._change_rngs(rv_out, rngs)
return rv_out
def rv_op(cls, rhos, sigma, init_dist, steps, ar_order, constant_term, size=None):
# Init dist should have shape (*size, ar_order)
if size is not None:
batch_size = size
else:
# In this case the size of the init_dist depends on the parameters shape
# The last dimension of rho and init_dist does not matter
batch_size = at.broadcast_shape(sigma, rhos[..., 0], init_dist[..., 0])
if init_dist.owner.op.ndim_supp == 0:
init_dist_size = (*batch_size, ar_order)
else:
# In this case the support dimension must cover for ar_order
init_dist_size = batch_size
init_dist = change_rv_size(init_dist, init_dist_size)
# Create OpFromGraph representing random draws form AR process
# Variables with underscore suffix are dummy inputs into the OpFromGraph
init_ = init_dist.type()
rhos_ = rhos.type()
sigma_ = sigma.type()
steps_ = steps.type()
rhos_bcast_shape_ = init_.shape
if constant_term:
# In this case init shape is one unit smaller than rhos in the last dimension
rhos_bcast_shape_ = (*rhos_bcast_shape_[:-1], rhos_bcast_shape_[-1] + 1)
rhos_bcast_ = at.broadcast_to(rhos_, rhos_bcast_shape_)
noise_rng = aesara.shared(np.random.default_rng())
def step(*args):
*prev_xs, reversed_rhos, sigma, rng = args
if constant_term:
mu = reversed_rhos[-1] + at.sum(prev_xs * reversed_rhos[:-1], axis=0)
else:
mu = at.sum(prev_xs * reversed_rhos, axis=0)
next_rng, new_x = Normal.dist(mu=mu, sigma=sigma, rng=rng).owner.outputs
return new_x, {rng: next_rng}
# We transpose inputs as scan iterates over first dimension
innov_, innov_updates_ = aesara.scan(
fn=step,
outputs_info=[{"initial": init_.T, "taps": range(-ar_order, 0)}],
non_sequences=[rhos_bcast_.T[::-1], sigma_.T, noise_rng],
n_steps=steps_,
strict=True,
)
(noise_next_rng,) = tuple(innov_updates_.values())
ar_ = at.concatenate([init_, innov_.T], axis=-1)
ar_op = AutoRegressiveRV(
inputs=[rhos_, sigma_, init_, steps_],
outputs=[noise_next_rng, ar_],
ar_order=ar_order,
constant_term=constant_term,
inline=True,
)
ar = ar_op(rhos, sigma, init_dist, steps)
return ar
def _resize_components(cls, size, *components):
if len(components) == 1:
# If we have a single component, we need to keep the length of the mixture
# axis intact, because that's what determines the number of mixture components
mix_axis = -components[0].owner.op.ndim_supp - 1
mix_size = components[0].shape[mix_axis]
size = tuple(size) + (mix_size, )
return [change_rv_size(component, size) for component in components]
def test_change_rv_size_default_update():
rng = aesara.shared(np.random.default_rng(0))
x = normal(rng=rng)
# Test that "traditional" default_update is updated
rng.default_update = x.owner.outputs[0]
new_x = change_rv_size(x, new_size=(2, ))
assert rng.default_update is not x.owner.outputs[0]
assert rng.default_update is new_x.owner.outputs[0]
# Test that "non-traditional" default_update is left unchanged
next_rng = aesara.shared(np.random.default_rng(1))
rng.default_update = next_rng
new_x = change_rv_size(x, new_size=(2, ))
assert rng.default_update is next_rng
# Test that default_update is not set if there was none before
del rng.default_update
new_x = change_rv_size(x, new_size=(2, ))
assert not hasattr(rng, "default_update")
def test_change_rv_size():
loc = at.as_tensor_variable([1, 2])
rv = normal(loc=loc)
assert rv.ndim == 1
assert tuple(rv.shape.eval()) == (2, )
with pytest.raises(ShapeError, match="must be ≤1-dimensional"):
change_rv_size(rv, new_size=[[2, 3]])
with pytest.raises(ShapeError, match="must be ≤1-dimensional"):
change_rv_size(rv, new_size=at.as_tensor_variable([[2, 3], [4, 5]]))
rv_new = change_rv_size(rv, new_size=(3, ), expand=True)
assert rv_new.ndim == 2
assert tuple(rv_new.shape.eval()) == (3, 2)
# Make sure that the shape used to determine the expanded size doesn't
# depend on the old `RandomVariable`.
rv_new_ancestors = set(ancestors((rv_new, )))
assert loc in rv_new_ancestors
assert rv not in rv_new_ancestors
rv_newer = change_rv_size(rv_new, new_size=(4, ), expand=True)
assert rv_newer.ndim == 3
assert tuple(rv_newer.shape.eval()) == (4, 3, 2)
# Make sure we avoid introducing a `Cast` by converting the new size before
# constructing the new `RandomVariable`
rv = normal(0, 1)
new_size = np.array([4, 3], dtype="int32")
rv_newer = change_rv_size(rv, new_size=new_size, expand=False)
assert rv_newer.ndim == 2
assert isinstance(rv_newer.owner.inputs[1], Constant)
assert tuple(rv_newer.shape.eval()) == (4, 3)
rv = normal(0, 1)
new_size = at.as_tensor(np.array([4, 3], dtype="int32"))
rv_newer = change_rv_size(rv, new_size=new_size, expand=True)
assert rv_newer.ndim == 2
assert tuple(rv_newer.shape.eval()) == (4, 3)
rv = normal(0, 1)
new_size = at.as_tensor(2, dtype="int32")
rv_newer = change_rv_size(rv, new_size=new_size, expand=True)
assert rv_newer.ndim == 1
assert tuple(rv_newer.shape.eval()) == (2, )
def make_node(self, rng, size, dtype, mu, sigma, init, steps):
steps = at.as_tensor_variable(steps)
if not steps.ndim == 0 or not steps.dtype.startswith("int"):
raise ValueError("steps must be an integer scalar (ndim=0).")
mu = at.as_tensor_variable(mu)
sigma = at.as_tensor_variable(sigma)
init = at.as_tensor_variable(init)
# Resize init distribution
size = normalize_size_param(size)
# If not explicit, size is determined by the shapes of mu, sigma, and init
init_size = size if not rv_size_is_none(size) else at.broadcast_shape(mu, sigma, init)
init = change_rv_size(init, init_size)
return super().make_node(rng, size, dtype, mu, sigma, init, steps)
def dist(
cls,
dist_params,
*,
shape: Optional[Shape] = None,
**kwargs,
) -> TensorVariable:
"""Creates a tensor variable corresponding to the `cls` distribution.
Parameters
----------
dist_params : array-like
The inputs to the `RandomVariable` `Op`.
shape : int, tuple, Variable, optional
A tuple of sizes for each dimension of the new RV.
An Ellipsis (...) may be inserted in the last position to short-hand refer to
all the dimensions that the RV would get if no shape/size/dims were passed at all.
**kwargs
Keyword arguments that will be forwarded to the Aesara RV Op.
Most prominently: ``size`` or ``dtype``.
Returns
-------
rv : TensorVariable
The created random variable tensor.
"""
if "testval" in kwargs:
kwargs.pop("testval")
warnings.warn(
"The `.dist(testval=...)` argument is deprecated and has no effect. "
"Initial values for sampling/optimization can be specified with `initval` in a modelcontext. "
"For using Aesara's test value features, you must assign the `.tag.test_value` yourself.",
FutureWarning,
stacklevel=2,
)
if "initval" in kwargs:
raise TypeError(
"Unexpected keyword argument `initval`. "
"This argument is not available for the `.dist()` API.")
if "dims" in kwargs:
raise NotImplementedError(
"The use of a `.dist(dims=...)` API is not supported.")
size = kwargs.pop("size", None)
if shape is not None and size is not None:
raise ValueError(
f"Passing both `shape` ({shape}) and `size` ({size}) is not supported!"
)
shape = convert_shape(shape)
size = convert_size(size)
create_size, ndim_expected, ndim_batch, ndim_supp = find_size(
shape=shape, size=size, ndim_supp=cls.rv_op.ndim_supp)
# Create the RV with a `size` right away.
# This is not necessarily the final result.
rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs)
# Replicate dimensions may be prepended via a shape with Ellipsis as the last element:
if shape is not None and Ellipsis in shape:
replicate_shape = cast(StrongShape, shape[:-1])
rv_out = change_rv_size(rv=rv_out,
new_size=replicate_shape,
expand=True)
rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")
rv_out.logcdf = _make_nice_attr_error("rv.logcdf(x)",
"pm.logcdf(rv, x)")
rv_out.random = _make_nice_attr_error("rv.random()", "rv.eval()")
return rv_out
def __new__(
cls,
name: str,
*args,
rng=None,
dims: Optional[Dims] = None,
initval=None,
observed=None,
total_size=None,
transform=UNSET,
**kwargs,
) -> TensorVariable:
"""Adds a tensor variable corresponding to a PyMC distribution to the current model.
Note that all remaining kwargs must be compatible with ``.dist()``
Parameters
----------
cls : type
A PyMC distribution.
name : str
Name for the new model variable.
rng : optional
Random number generator to use with the RandomVariable.
dims : tuple, optional
A tuple of dimension names known to the model.
initval : optional
Numeric or symbolic untransformed initial value of matching shape,
or one of the following initial value strategies: "moment", "prior".
Depending on the sampler's settings, a random jitter may be added to numeric, symbolic
or moment-based initial values in the transformed space.
observed : optional
Observed data to be passed when registering the random variable in the model.
See ``Model.register_rv``.
total_size : float, optional
See ``Model.register_rv``.
transform : optional
See ``Model.register_rv``.
**kwargs
Keyword arguments that will be forwarded to ``.dist()`` or the Aesara RV Op.
Most prominently: ``shape`` for ``.dist()`` or ``dtype`` for the Op.
Returns
-------
rv : TensorVariable
The created random variable tensor, registered in the Model.
"""
try:
from pymc.model import Model
model = Model.get_context()
except TypeError:
raise TypeError("No model on context stack, which is needed to "
"instantiate distributions. Add variable inside "
"a 'with model:' block, or use the '.dist' syntax "
"for a standalone distribution.")
if "testval" in kwargs:
initval = kwargs.pop("testval")
warnings.warn(
"The `testval` argument is deprecated; use `initval`.",
FutureWarning,
stacklevel=2,
)
if not isinstance(name, string_types):
raise TypeError(f"Name needs to be a string but got: {name}")
# Create the RV and process dims and observed to determine
# a shape by which the created RV may need to be resized.
rv_out, dims, observed, resize_shape = _make_rv_and_resize_shape(
cls=cls,
dims=dims,
model=model,
observed=observed,
args=args,
**kwargs)
if resize_shape:
# A batch size was specified through `dims`, or implied by `observed`.
rv_out = change_rv_size(rv=rv_out,
new_size=resize_shape,
expand=True)
rv_out = model.register_rv(
rv_out,
name,
observed,
total_size,
dims=dims,
transform=transform,
initval=initval,
)
# add in pretty-printing support
rv_out.str_repr = types.MethodType(str_for_dist, rv_out)
rv_out._repr_latex_ = types.MethodType(
functools.partial(str_for_dist, formatting="latex"), rv_out)
rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")
rv_out.logcdf = _make_nice_attr_error("rv.logcdf(x)",
"pm.logcdf(rv, x)")
rv_out.random = _make_nice_attr_error("rv.random()", "rv.eval()")
return rv_out
def change_size(cls, rv, new_size, expand=False):
dist = rv.tag.dist
lower = rv.tag.lower
upper = rv.tag.upper
new_dist = change_rv_size(dist, new_size, expand=expand)
return cls.rv_op(new_dist, lower, upper)
def maybe_resize(
rv_out,
rv_op,
dist_params,
ndim_expected,
ndim_batch,
ndim_supp,
shape,
size,
**kwargs,
):
"""Resize a distribution if necessary.
Parameters
----------
rv_out : RandomVariable
The RandomVariable to be resized if necessary
rv_op : RandomVariable.__class__
The RandomVariable class to recreate it
dist_params : dict
Input parameters to recreate the RandomVariable
ndim_expected : int
Number of dimensions expected after distribution was created
ndim_batch : int
Number of batch dimensions
ndim_supp : int
The support dimension of the distribution.
0 if a univariate distribution, 1 if a multivariate distribution.
shape : tuple
A tuple specifying the final shape of a distribution
size : tuple
A tuple specifying the size of a distribution
Returns
-------
rv_out : int
The size argument to be passed to the distribution
"""
ndim_actual = rv_out.ndim
ndims_unexpected = ndim_actual != ndim_expected
if shape is not None and ndims_unexpected:
if Ellipsis in shape:
# Resize and we're done!
rv_out = change_rv_size(rv_var=rv_out,
new_size=shape[:-1],
expand=True)
else:
# This is rare, but happens, for example, with MvNormal(np.ones((2, 3)), np.eye(3), shape=(2, 3)).
# Recreate the RV without passing `size` to created it with just the implied dimensions.
rv_out = rv_op(*dist_params, size=None, **kwargs)
# Now resize by any remaining "extra" dimensions that were not implied from support and parameters
if rv_out.ndim < ndim_expected:
expand_shape = shape[:ndim_expected - rv_out.ndim]
rv_out = change_rv_size(rv_var=rv_out,
new_size=expand_shape,
expand=True)
if not rv_out.ndim == ndim_expected:
raise ShapeError(
f"Failed to create the RV with the expected dimensionality. "
f"This indicates a severe problem. Please open an issue.",
actual=ndim_actual,
expected=ndim_batch + ndim_supp,
)
# Warn about the edge cases where the RV Op creates more dimensions than
# it should based on `size` and `RVOp.ndim_supp`.
if size is not None and ndims_unexpected:
warnings.warn(
f"You may have expected a ({len(tuple(size))}+{ndim_supp})-dimensional RV, but the resulting RV will be {ndim_actual}-dimensional."
' To silence this warning use `warnings.simplefilter("ignore", pm.ShapeWarning)`.',
ShapeWarning,
)
return rv_out
def __new__(
cls,
name: str,
*args,
rng=None,
dims: Optional[Dims] = None,
initval=None,
observed=None,
total_size=None,
transform=UNSET,
**kwargs,
) -> RandomVariable:
"""Adds a RandomVariable corresponding to a PyMC distribution to the current model.
Note that all remaining kwargs must be compatible with ``.dist()``
Parameters
----------
cls : type
A PyMC distribution.
name : str
Name for the new model variable.
rng : optional
Random number generator to use with the RandomVariable.
dims : tuple, optional
A tuple of dimension names known to the model.
initval : optional
Numeric or symbolic untransformed initial value of matching shape,
or one of the following initial value strategies: "moment", "prior".
Depending on the sampler's settings, a random jitter may be added to numeric, symbolic
or moment-based initial values in the transformed space.
observed : optional
Observed data to be passed when registering the random variable in the model.
See ``Model.register_rv``.
total_size : float, optional
See ``Model.register_rv``.
transform : optional
See ``Model.register_rv``.
**kwargs
Keyword arguments that will be forwarded to ``.dist()``.
Most prominently: ``shape`` and ``size``
Returns
-------
rv : RandomVariable
The created RV, registered in the Model.
"""
try:
from pymc.model import Model
model = Model.get_context()
except TypeError:
raise TypeError("No model on context stack, which is needed to "
"instantiate distributions. Add variable inside "
"a 'with model:' block, or use the '.dist' syntax "
"for a standalone distribution.")
if "testval" in kwargs:
initval = kwargs.pop("testval")
warnings.warn(
"The `testval` argument is deprecated; use `initval`.",
DeprecationWarning,
stacklevel=2,
)
if not isinstance(name, string_types):
raise TypeError(f"Name needs to be a string but got: {name}")
if rng is None:
rng = model.next_rng()
if dims is not None and "shape" in kwargs:
raise ValueError(
f"Passing both `dims` ({dims}) and `shape` ({kwargs['shape']}) is not supported!"
)
if dims is not None and "size" in kwargs:
raise ValueError(
f"Passing both `dims` ({dims}) and `size` ({kwargs['size']}) is not supported!"
)
dims = convert_dims(dims)
# Create the RV without dims information, because that's not something tracked at the Aesara level.
# If necessary we'll later replicate to a different size implied by already known dims.
rv_out = cls.dist(*args, rng=rng, **kwargs)
ndim_actual = rv_out.ndim
resize_shape = None
# `dims` are only available with this API, because `.dist()` can be used
# without a modelcontext and dims are not tracked at the Aesara level.
if dims is not None:
ndim_resize, resize_shape, dims = resize_from_dims(
dims, ndim_actual, model)
elif observed is not None:
ndim_resize, resize_shape, observed = resize_from_observed(
observed, ndim_actual)
if resize_shape:
# A batch size was specified through `dims`, or implied by `observed`.
rv_out = change_rv_size(rv_var=rv_out,
new_size=resize_shape,
expand=True)
rv_out = model.register_rv(
rv_out,
name,
observed,
total_size,
dims=dims,
transform=transform,
initval=initval,
)
# add in pretty-printing support
rv_out.str_repr = types.MethodType(str_for_dist, rv_out)
rv_out._repr_latex_ = types.MethodType(
functools.partial(str_for_dist, formatting="latex"), rv_out)
rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")
rv_out.logcdf = _make_nice_attr_error("rv.logcdf(x)",
"pm.logcdf(rv, x)")
rv_out.random = _make_nice_attr_error("rv.random()", "rv.eval()")
return rv_out
def dist(
cls,
dist_params,
*,
shape: Optional[Shape] = None,
size: Optional[Size] = None,
**kwargs,
) -> RandomVariable:
"""Creates a RandomVariable corresponding to the `cls` distribution.
Parameters
----------
dist_params : array-like
The inputs to the `RandomVariable` `Op`.
shape : int, tuple, Variable, optional
A tuple of sizes for each dimension of the new RV.
An Ellipsis (...) may be inserted in the last position to short-hand refer to
all the dimensions that the RV would get if no shape/size/dims were passed at all.
size : int, tuple, Variable, optional
For creating the RV like in Aesara/NumPy.
Returns
-------
rv : RandomVariable
The created RV.
"""
if "testval" in kwargs:
kwargs.pop("testval")
warnings.warn(
"The `.dist(testval=...)` argument is deprecated and has no effect. "
"Initial values for sampling/optimization can be specified with `initval` in a modelcontext. "
"For using Aesara's test value features, you must assign the `.tag.test_value` yourself.",
FutureWarning,
stacklevel=2,
)
if "initval" in kwargs:
raise TypeError(
"Unexpected keyword argument `initval`. "
"This argument is not available for the `.dist()` API.")
if "dims" in kwargs:
raise NotImplementedError(
"The use of a `.dist(dims=...)` API is not supported.")
if shape is not None and size is not None:
raise ValueError(
f"Passing both `shape` ({shape}) and `size` ({size}) is not supported!"
)
shape = convert_shape(shape)
size = convert_size(size)
create_size, ndim_expected, ndim_batch, ndim_supp = find_size(
shape=shape, size=size, ndim_supp=cls.rv_op.ndim_supp)
# Create the RV with a `size` right away.
# This is not necessarily the final result.
rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs)
# Replicate dimensions may be prepended via a shape with Ellipsis as the last element:
if shape is not None and Ellipsis in shape:
replicate_shape = cast(StrongShape, shape[:-1])
rv_out = change_rv_size(rv_var=rv_out,
new_size=replicate_shape,
expand=True)
rng = kwargs.pop("rng", None)
if (rv_out.owner and isinstance(rv_out.owner.op, RandomVariable)
and isinstance(rng, RandomStateSharedVariable)
and not getattr(rng, "default_update", None)):
# This tells `aesara.function` that the shared RNG variable
# is mutable, which--in turn--tells the `FunctionGraph`
# `Supervisor` feature to allow in-place updates on the variable.
# Without it, the `RandomVariable`s could not be optimized to allow
# in-place RNG updates, forcing all sample results from compiled
# functions to be the same on repeated evaluations.
new_rng = rv_out.owner.outputs[0]
rv_out.update = (rng, new_rng)
rng.default_update = new_rng
rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")
rv_out.logcdf = _make_nice_attr_error("rv.logcdf(x)",
"pm.logcdf(rv, x)")
rv_out.random = _make_nice_attr_error("rv.random()", "rv.eval()")
return rv_out
