def local_rv_size_lift(fgraph, node):
"""Lift the ``size`` parameter in a ``RandomVariable``.
In other words, this will broadcast the distribution parameters by adding
the extra dimensions implied by the ``size`` parameter, and remove the
``size`` parameter in the process.
For example, ``normal(0, 1, size=(1, 2))`` becomes
``normal([[0, 0]], [[1, 1]], size=())``.
"""
if not isinstance(node.op, RandomVariable):
return
rng, size, dtype, *dist_params = node.inputs
dist_params = broadcast_params(dist_params, node.op.ndims_params)
if get_vector_length(size) > 0:
dist_params = [
broadcast_to(p,
(tuple(size) +
tuple(p.shape)) if node.op.ndim_supp > 0 else size)
for p in dist_params
]
else:
return
new_node = node.op.make_node(rng, None, dtype, *dist_params)
if config.compute_test_value != "off":
compute_test_value(new_node)
return new_node.outputs
def lift_rv_shapes(node):
"""Lift `RandomVariable`'s shape-related parameters.
In other words, this will broadcast the distribution parameters and
extra dimensions added by the `size` parameter.
For example, ``normal([0.0, 1.0], 5.0, size=(3, 2))`` becomes
``normal([[0., 1.], [0., 1.], [0., 1.]], [[5., 5.], [5., 5.], [5., 5.]])``.
"""
if not isinstance(node.op, RandomVariable):
return False
rng, size, dtype, *dist_params = node.inputs
dist_params = broadcast_params(dist_params, node.op.ndims_params)
if get_vector_length(size) > 0:
dist_params = [
broadcast_to(p,
(tuple(size) +
tuple(p.shape)) if node.op.ndim_supp > 0 else size)
for p in dist_params
]
new_node = node.op.make_node(rng, None, dtype, *dist_params)
if config.compute_test_value != "off":
compute_test_value(new_node)
return new_node
def numba_funcify_CategoricalRV(op, node, **kwargs):
out_dtype = node.outputs[1].type.numpy_dtype
ind_shape_len = node.inputs[3].type.ndim - 1
neg_ind_shape_len = -ind_shape_len
size_len = int(get_vector_length(node.inputs[1]))
@numba_basic.numba_njit
def categorical_rv(rng, size, dtype, p):
size_tpl = numba_ndarray.to_fixed_tuple(size, size_len)
ind_shape = p.shape[:-1]
if ind_shape_len > 0:
if size_len > 0 and size_tpl[neg_ind_shape_len:] != ind_shape:
raise ValueError("Parameters shape and size do not match.")
samples_shape = size_tpl[:neg_ind_shape_len] + ind_shape
p_bcast = np.broadcast_to(p,
size_tpl[:neg_ind_shape_len] + p.shape)
else:
samples_shape = size_tpl
p_bcast = p
unif_samples = np.random.uniform(0, 1, samples_shape)
res = np.empty(samples_shape, dtype=out_dtype)
for idx in np.ndindex(*samples_shape):
res[idx] = np.searchsorted(np.cumsum(p_bcast[idx]),
unif_samples[idx])
return (rng, res)
return categorical_rv
def infer_shape(self, fgraph, node, input_shapes):
_, size, _, *dist_params = node.inputs
_, size_shape, _, *param_shapes = input_shapes
try:
size_len = get_vector_length(size)
except ValueError:
size_len = get_scalar_constant_value(size_shape[0])
size = tuple(size[n] for n in range(size_len))
shape = self._infer_shape(size, dist_params, param_shapes=param_shapes)
return [None, [s for s in shape]]
def reshape(x, newshape, ndim=None):
if ndim is None:
newshape = aet.as_tensor_variable(newshape)
if newshape.ndim != 1:
raise TypeError(
"New shape in reshape must be a vector or a list/tuple of"
f" scalar. Got {newshape} after conversion to a vector.")
try:
ndim = aet.get_vector_length(newshape)
except ValueError:
raise ValueError(
f"The length of the provided shape ({newshape}) cannot "
"be automatically determined, so Aesara is not able "
"to know what the number of dimensions of the reshaped "
"variable will be. You can provide the 'ndim' keyword "
"argument to 'reshape' to avoid this problem.")
op = Reshape(ndim)
rval = op(x, newshape)
return rval
def make_node(self, indices, dims):
indices = aet.as_tensor_variable(indices)
dims = aet.as_tensor_variable(dims)
if indices.dtype not in int_dtypes:
raise TypeError(
f"'{indices.dtype}' object cannot be interpreted as an index"
)
if dims.dtype not in int_dtypes:
raise TypeError(f"'{dims.dtype}' object cannot be interpreted as an index")
if dims.ndim != 1:
raise TypeError("dims must be a 1D array")
return Apply(
self,
[indices, dims],
[
TensorType(dtype="int64", broadcastable=(False,) * indices.ndim)()
for i in range(aet.get_vector_length(dims))
],
)
def numba_funcify_DirichletRV(op, node, **kwargs):
out_dtype = node.outputs[1].type.numpy_dtype
alphas_ndim = node.inputs[3].type.ndim
neg_ind_shape_len = -alphas_ndim + 1
size_len = int(get_vector_length(node.inputs[1]))
if alphas_ndim > 1:
@numba_basic.numba_njit
def dirichlet_rv(rng, size, dtype, alphas):
if size_len > 0:
size_tpl = numba_ndarray.to_fixed_tuple(size, size_len)
if (0 < alphas.ndim - 1 <= len(size_tpl)
and size_tpl[neg_ind_shape_len:] != alphas.shape[:-1]):
raise ValueError("Parameters shape and size do not match.")
samples_shape = size_tpl + alphas.shape[-1:]
else:
samples_shape = alphas.shape
res = np.empty(samples_shape, dtype=out_dtype)
alphas_bcast = np.broadcast_to(alphas, samples_shape)
for index in np.ndindex(*samples_shape[:-1]):
res[index] = np.random.dirichlet(alphas_bcast[index])
return (rng, res)
else:
@numba_basic.numba_njit
def dirichlet_rv(rng, size, dtype, alphas):
size = numba_ndarray.to_fixed_tuple(size, size_len)
return (rng, np.random.dirichlet(alphas, size))
return dirichlet_rv
def local_dimshuffle_rv_lift(fgraph, node):
"""Lift `DimShuffle`s through `RandomVariable` `Op`s.
For example, ``normal(mu, std).T == normal(mu.T, std.T)``.
The basic idea behind this optimization is that we need to separate the
`DimShuffle`ing into independent `DimShuffle`s that each occur in two
distinct sub-spaces: the parameters and ``size`` (i.e. replications)
sub-spaces.
If a `DimShuffle` exchanges dimensions across those two sub-spaces, then we
don't do anything.
Otherwise, if the `DimShuffle` only exchanges dimensions within each of
those sub-spaces, we can break it apart and apply the parameter-space
`DimShuffle` to the `RandomVariable`'s distribution parameters, and the
apply the replications-space `DimShuffle` to the `RandomVariable`'s``size``
tuple. The latter is a particularly simple rearranging of a tuple, but the
former requires a little more work.
"""
ds_op = node.op
if not isinstance(ds_op, DimShuffle):
return False
base_rv = node.inputs[0]
rv_node = base_rv.owner
if not (
rv_node and isinstance(rv_node.op, RandomVariable) and rv_node.op.ndim_supp == 0
):
return False
# If no one else is using the underlying `RandomVariable`, then we can
# do this; otherwise, the graph would be internally inconsistent.
if not all(
(n == node or isinstance(n.op, Shape)) for n, i in fgraph.clients[base_rv]
):
return False
rv_op = rv_node.op
rng, size, dtype, *dist_params = rv_node.inputs
# We need to know the dimensions that were *not* added by the `size`
# parameter (i.e. the dimensions corresponding to independent variates with
# different parameter values)
num_ind_dims = None
if len(dist_params) == 1:
num_ind_dims = dist_params[0].ndim
else:
# When there is more than one distribution parameter, assume that all
# of them will broadcast to the maximum number of dimensions
num_ind_dims = max(d.ndim for d in dist_params)
# If the indices in `ds_new_order` are entirely within the replication
# indices group or the independent variates indices group, then we can apply
# this optimization.
ds_new_order = ds_op.new_order
# Create a map from old index order to new/`DimShuffled` index order
dim_orders = [(n, d) for n, d in enumerate(ds_new_order) if isinstance(d, int)]
# Find the index at which the replications/independents split occurs
reps_ind_split_idx = len(dim_orders) - (num_ind_dims + rv_op.ndim_supp)
ds_reps_new_dims = dim_orders[:reps_ind_split_idx]
ds_ind_new_dims = dim_orders[reps_ind_split_idx:]
ds_only_in_ind = ds_ind_new_dims and all(
d >= reps_ind_split_idx for n, d in ds_ind_new_dims
)
if ds_only_in_ind:
# Update the `size` array to reflect the `DimShuffle`d dimensions,
# since the trailing dimensions in `size` represent the independent
# variates dimensions (for univariate distributions, at least)
new_size = (
[constant(1, dtype="int64") if o == "x" else size[o] for o in ds_new_order]
if get_vector_length(size) > 0
else size
)
# Compute the new axes parameter(s) for the `DimShuffle` that will be
# applied to the `RandomVariable` parameters (they need to be offset)
rv_params_new_order = [
d - reps_ind_split_idx if isinstance(d, int) else d
for d in ds_new_order[ds_ind_new_dims[0][0] :]
]
# Lift the `DimShuffle`s into the parameters
# NOTE: The parameters might not be broadcasted against each other, so
# we can only apply the parts of the `DimShuffle` that are relevant.
new_dist_params = []
for d in dist_params:
if d.ndim < len(ds_ind_new_dims):
_rv_params_new_order = [
o
for o in rv_params_new_order
if (isinstance(o, int) and o < d.ndim) or o == "x"
]
else:
_rv_params_new_order = rv_params_new_order
new_dist_params.append(
type(ds_op)(d.type.broadcastable, _rv_params_new_order)(d)
)
new_node = rv_op.make_node(rng, new_size, dtype, *new_dist_params)
if config.compute_test_value != "off":
compute_test_value(new_node)
return [new_node.outputs[1]]
ds_only_in_reps = ds_reps_new_dims and all(
d < reps_ind_split_idx for n, d in ds_reps_new_dims
)
if ds_only_in_reps:
# Update the `size` array to reflect the `DimShuffle`d dimensions.
# There should be no need to `DimShuffle` now.
new_size = [
constant(1, dtype="int64") if o == "x" else size[o] for o in ds_new_order
]
new_node = rv_op.make_node(rng, new_size, dtype, *dist_params)
if config.compute_test_value != "off":
compute_test_value(new_node)
return [new_node.outputs[1]]
return False
def _infer_shape(
self,
size: Tuple[TensorVariable],
dist_params: List[TensorVariable],
param_shapes: Optional[List[Tuple[TensorVariable]]] = None,
) -> Tuple[ScalarVariable]:
"""Compute the output shape given the size and distribution parameters.
Parameters
----------
size
The size parameter specified for this `RandomVariable`.
dist_params
The symbolic parameter for this `RandomVariable`'s distribution.
param_shapes
The shapes of the `dist_params` as given by `ShapeFeature`'s
via `Op.infer_shape`'s `input_shapes` argument. This parameter's
values are essentially more accurate versions of ``[d.shape for d
in dist_params]``.
"""
size_len = get_vector_length(size)
if self.ndim_supp == 0 and size_len > 0:
# In this case, we have a univariate distribution with a non-empty
# `size` parameter, which means that the `size` parameter
# completely determines the shape of the random variable. More
# importantly, the `size` parameter may be the only correct source
# of information for the output shape, in that we would be misled
# by the `dist_params` if we tried to infer the relevant parts of
# the output shape from those.
return size
# Broadcast the parameters
param_shapes = params_broadcast_shapes(
param_shapes or [shape_tuple(p) for p in dist_params],
self.ndims_params)
def slice_ind_dims(p, ps, n):
shape = tuple(ps)
if n == 0:
return (p, shape)
ind_slice = (slice(None), ) * (p.ndim - n) + (0, ) * n
ind_shape = [
s if b is False else constant(1, "int64")
for s, b in zip(shape[:-n], p.broadcastable[:-n])
]
return (
p[ind_slice],
ind_shape,
)
# These are versions of our actual parameters with the anticipated
# dimensions (i.e. support dimensions) removed so that only the
# independent variate dimensions are left.
params_ind_slice = tuple(
slice_ind_dims(p, ps, n)
for p, ps, n in zip(dist_params, param_shapes, self.ndims_params))
if len(params_ind_slice) == 1:
ind_param, ind_shape = params_ind_slice[0]
ndim_ind = len(ind_shape)
shape_ind = ind_shape
elif len(params_ind_slice) > 1:
# If there are multiple parameters, the dimensions of their
# independent variates should broadcast together.
p_slices, p_shapes = zip(*params_ind_slice)
shape_ind = aesara.tensor.extra_ops.broadcast_shape_iter(
p_shapes, arrays_are_shapes=True)
ndim_ind = len(shape_ind)
else:
ndim_ind = 0
if self.ndim_supp == 0:
shape_supp = tuple()
shape_reps = tuple(size)
if ndim_ind > 0:
shape_reps = shape_reps[:-ndim_ind]
ndim_reps = len(shape_reps)
else:
shape_supp = self._shape_from_params(
dist_params,
param_shapes=param_shapes,
)
ndim_reps = size_len
shape_reps = size
ndim_shape = self.ndim_supp + ndim_ind + ndim_reps
if ndim_shape == 0:
shape = constant([], dtype="int64")
else:
shape = tuple(shape_reps) + tuple(shape_ind) + tuple(shape_supp)
# if shape is None:
# raise ShapeError()
return shape
def make_numba_random_fn(node, np_random_func):
"""Create Numba implementations for existing Numba-supported ``np.random`` functions.
The functions generated here add parameter broadcasting and the ``size``
argument to the Numba-supported scalar ``np.random`` functions.
"""
tuple_size = int(get_vector_length(node.inputs[1]))
size_dims = tuple_size - max(i.ndim for i in node.inputs[3:])
# Make a broadcast-capable version of the Numba supported scalar sampling
# function
bcast_fn_name = f"aesara_random_{get_name_for_object(np_random_func)}"
sized_fn_name = "sized_random_variable"
unique_names = unique_name_generator(
[
bcast_fn_name,
sized_fn_name,
"np",
"np_random_func",
"numba_vectorize",
"to_fixed_tuple",
"tuple_size",
"size_dims",
"rng",
"size",
"dtype",
],
suffix_sep="_",
)
bcast_fn_input_names = ", ".join(
[unique_names(i, force_unique=True) for i in node.inputs[3:]])
bcast_fn_global_env = {
"np_random_func": np_random_func,
"numba_vectorize": numba.vectorize,
}
bcast_fn_src = f"""
@numba_vectorize
def {bcast_fn_name}({bcast_fn_input_names}):
return np_random_func({bcast_fn_input_names})
"""
bcast_fn = compile_function_src(bcast_fn_src, bcast_fn_name,
bcast_fn_global_env)
random_fn_input_names = ", ".join(
["rng", "size", "dtype"] + [unique_names(i) for i in node.inputs[3:]])
# Now, create a Numba JITable function that implements the `size` parameter
out_dtype = node.outputs[1].type.numpy_dtype
random_fn_global_env = {
bcast_fn_name: bcast_fn,
"out_dtype": out_dtype,
}
if tuple_size > 0:
random_fn_body = dedent(f"""
size = to_fixed_tuple(size, tuple_size)
data = np.empty(size, dtype=out_dtype)
for i in np.ndindex(size[:size_dims]):
data[i] = {bcast_fn_name}({bcast_fn_input_names})
""")
random_fn_global_env.update({
"np": np,
"to_fixed_tuple": numba_ndarray.to_fixed_tuple,
"tuple_size": tuple_size,
"size_dims": size_dims,
})
else:
random_fn_body = f"""data = {bcast_fn_name}({bcast_fn_input_names})"""
sized_fn_src = dedent(f"""
def {sized_fn_name}({random_fn_input_names}):
{indent(random_fn_body, " " * 4)}
return (rng, data)
""")
random_fn = compile_function_src(sized_fn_src, sized_fn_name,
random_fn_global_env)
random_fn = numba.njit(random_fn)
return random_fn
def _infer_shape(
self,
size: TensorVariable,
dist_params: Sequence[TensorVariable],
param_shapes: Optional[Sequence[Tuple[Variable, ...]]] = None,
) -> Union[TensorVariable, Tuple[ScalarVariable, ...]]:
"""Compute the output shape given the size and distribution parameters.
Parameters
----------
size
The size parameter specified for this `RandomVariable`.
dist_params
The symbolic parameter for this `RandomVariable`'s distribution.
param_shapes
The shapes of the `dist_params` as given by `ShapeFeature`'s
via `Op.infer_shape`'s `input_shapes` argument. This parameter's
values are essentially more accurate versions of ``[d.shape for d
in dist_params]``.
"""
size_len = get_vector_length(size)
if size_len > 0:
if self.ndim_supp == 0:
return size
else:
supp_shape = self._supp_shape_from_params(
dist_params, param_shapes=param_shapes)
return tuple(size) + tuple(supp_shape)
# Broadcast the parameters
param_shapes = params_broadcast_shapes(
param_shapes or [shape_tuple(p) for p in dist_params],
self.ndims_params)
def slice_ind_dims(p, ps, n):
shape = tuple(ps)
if n == 0:
return (p, shape)
ind_slice = (slice(None), ) * (p.ndim - n) + (0, ) * n
ind_shape = [
s if b is False else constant(1, "int64")
for s, b in zip(shape[:-n], p.broadcastable[:-n])
]
return (
p[ind_slice],
ind_shape,
)
# These are versions of our actual parameters with the anticipated
# dimensions (i.e. support dimensions) removed so that only the
# independent variate dimensions are left.
params_ind_slice = tuple(
slice_ind_dims(p, ps, n)
for p, ps, n in zip(dist_params, param_shapes, self.ndims_params))
if len(params_ind_slice) == 1:
_, shape_ind = params_ind_slice[0]
elif len(params_ind_slice) > 1:
# If there are multiple parameters, the dimensions of their
# independent variates should broadcast together.
p_slices, p_shapes = zip(*params_ind_slice)
shape_ind = aesara.tensor.extra_ops.broadcast_shape_iter(
p_shapes, arrays_are_shapes=True)
else:
# Distribution has no parameters
shape_ind = ()
if self.ndim_supp == 0:
shape_supp = ()
else:
shape_supp = self._supp_shape_from_params(
dist_params,
param_shapes=param_shapes,
)
shape = tuple(shape_ind) + tuple(shape_supp)
if not shape:
shape = constant([], dtype="int64")
return shape
