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