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