def test_inputs_owners(self): # Test a graph where the inputs have owners r1, r5 = MyVariable(1), MyVariable(5) o = MyOp.make_node(r1, r1) r2b = o.outputs[0] o2 = MyOp.make_node(r2b, r2b) all = io_toposort([r2b], o2.outputs) assert all == [o2] o2 = MyOp.make_node(r2b, r5) all = io_toposort([r2b], o2.outputs) assert all == [o2]
def import_node( self, apply_node: Apply, check: bool = True, reason: str = None, import_missing: bool = False, ) -> NoReturn: """Recursively import everything between an `Apply` node and the `FunctionGraph`'s outputs. Parameters: ---------- apply_node : aesara.graph.basic.Apply The node to be imported. check : bool Check that the inputs for the imported nodes are also present in the `FunctionGraph`. reason : str The name of the optimization or operation in progress. import_missing : bool Add missing inputs instead of raising an exception. """ # We import the nodes in topological order. We only are interested in # new nodes, so we use all variables we know of as if they were the # input set. (The functions in the graph module only use the input set # to know where to stop going down.) new_nodes = io_toposort(self.variables, apply_node.outputs) if check: for node in new_nodes: for var in node.inputs: if (var.owner is None and not isinstance(var, Constant) and var not in self.inputs): if import_missing: self.add_input(var) else: error_msg = ( f"Input {node.inputs.index(var)} ({var})" " of the graph (indices start " f"from 0), used to compute {node}, was not " "provided and not given a value. Use the " "Aesara flag exception_verbosity='high', " "for more information on this error.") raise MissingInputError(error_msg, variable=var) for node in new_nodes: assert node not in self.apply_nodes self.setup_node(node) self.apply_nodes.add(node) if not hasattr(node.tag, "imported_by"): node.tag.imported_by = [] node.tag.imported_by.append(str(reason)) for output in node.outputs: self.setup_var(output) self.variables.add(output) for i, input in enumerate(node.inputs): if input not in self.variables: self.setup_var(input) self.variables.add(input) self.add_client(input, (node, i)) self.execute_callbacks("on_import", node, reason)
def on_detach(self, fgraph): """ Should remove any dynamically added functionality that it installed into the function_graph """ for node in io_toposort(fgraph.inputs, fgraph.outputs): self.on_prune(fgraph, node, "Bookkeeper.detach")
def test_outputs_clients(self): # Test when outputs have clients r1, r2, r4 = MyVariable(1), MyVariable(2), MyVariable(4) o0 = MyOp.make_node(r1, r2) MyOp.make_node(o0.outputs[0], r4) all = io_toposort([], o0.outputs) assert all == [o0]
def test_not_connected(self): # Test a graph which is not connected r1, r2, r3, r4 = MyVariable(1), MyVariable(2), MyVariable(3), MyVariable(4) o0 = MyOp.make_node(r1, r2) o1 = MyOp.make_node(r3, r4) all = io_toposort([r1, r2, r3, r4], o0.outputs + o1.outputs) assert all == [o1, o0] or all == [o0, o1]
def toposort(self) -> List[Apply]: """Toposort Return an ordering of the graph's Apply nodes such that * All the nodes of the inputs of a node are before that node. * Satisfies the orderings provided by each feature that has an 'orderings' method. If a feature has an 'orderings' method, it will be called with this FunctionGraph as sole argument. It should return a dictionary of `{node: predecessors}` where predecessors is a list of nodes that should be computed before the key node. """ if len(self.apply_nodes) < 2: # optimization # when there are 0 or 1 nodes, no sorting is necessary # This special case happens a lot because the OpWiseCLinker # produces 1-element graphs. return list(self.apply_nodes) fg = self ords = self.orderings() order = io_toposort(fg.inputs, fg.outputs, ords) return order
def test_simple(self): # Test a simple graph r1, r2, r5 = MyVariable(1), MyVariable(2), MyVariable(5) o = MyOp(r1, r2) o.name = "o1" o2 = MyOp(o, r5) o2.name = "o2" clients = {} res = general_toposort([o2], prenode, clients=clients) assert clients == { o2.owner: [o2], o: [o2.owner], r5: [o2.owner], o.owner: [o], r1: [o.owner], r2: [o.owner], } assert res == [r5, r2, r1, o.owner, o, o2.owner, o2] with pytest.raises(ValueError): general_toposort( [o2], prenode, compute_deps_cache=lambda x: None, deps_cache=None ) res = io_toposort([r5], [o2]) assert res == [o.owner, o2.owner]
def test_io_chain(self): # Test inputs and outputs mixed together in a chain graph r1, r2 = MyVariable(1), MyVariable(2) o0 = MyOp.make_node(r1, r2) o1 = MyOp.make_node(o0.outputs[0], r1) all = io_toposort([r1, o0.outputs[0]], [o0.outputs[0], o1.outputs[0]]) assert all == [o1]
def on_attach(self, fgraph): """ Called by FunctionGraph.attach_feature, the method that attaches the feature to the FunctionGraph. Since this is called after the FunctionGraph is initially populated, this is where you should run checks on the initial contents of the FunctionGraph. """ for node in io_toposort(fgraph.inputs, fgraph.outputs): self.on_import(fgraph, node, "on_attach")
def test_dependence(): dependence = make_dependence_cmp() x = matrix("x") y = dot(x * 2, x + 1) nodes = io_toposort([x], [y]) for a, b in zip(nodes[:-1], nodes[1:]): assert dependence(a, b) <= 0
def process_graph(self, inputs, outputs, updates=None, display_inputs=False): if updates is None: updates = {} if not isinstance(inputs, (list, tuple)): inputs = [inputs] if not isinstance(outputs, (list, tuple)): outputs = [outputs] current = None if display_inputs: strings = [ (0, "inputs: " + ", ".join(map(str, list(inputs) + updates.keys()))) ] else: strings = [] pprinter = self.clone_assign( lambda pstate, r: r.name is not None and r is not current, leaf_printer) inv_updates = {b: a for (a, b) in updates.items()} i = 1 for node in io_toposort( list(inputs) + updates.keys(), list(outputs) + updates.values()): for output in node.outputs: if output in inv_updates: name = str(inv_updates[output]) strings.append( (i + 1000, f"{name} <- {pprinter.process(output)}")) i += 1 if output.name is not None or output in outputs: if output.name is None: name = "out[%i]" % outputs.index(output) else: name = output.name # backport # name = 'out[%i]' % outputs.index(output) if output.name # is None else output.name current = output try: idx = 2000 + outputs.index(output) except ValueError: idx = i if len(outputs) == 1 and outputs[0] is output: strings.append( (idx, f"return {pprinter.process(output)}")) else: strings.append( (idx, f"{name} = {pprinter.process(output)}")) i += 1 strings.sort() return "\n".join(s[1] for s in strings)
def toposort(self) -> List[Apply]: r"""Return a toposorted list of the nodes. Return an ordering of the graph's :class:`Apply` nodes such that: * all the nodes of the inputs of a node are before that node, and * they satisfy the additional orderings provided by :meth:`FunctionGraph.orderings`. """ if len(self.apply_nodes) < 2: # No sorting is necessary return list(self.apply_nodes) return io_toposort(self.inputs, self.outputs, self.orderings())
def joint_logpt( var: Union[TensorVariable, List[TensorVariable]], rv_values: Optional[Union[TensorVariable, Dict[TensorVariable, TensorVariable]]] = None, *, jacobian: bool = True, scaling: bool = True, transformed: bool = True, sum: bool = True, **kwargs, ) -> Union[TensorVariable, List[TensorVariable]]: """Create a measure-space (i.e. log-likelihood) graph for a random variable or a list of random variables at a given point. The input `var` determines which log-likelihood graph is used and `rv_value` is that graph's input parameter. For example, if `var` is the output of a ``NormalRV`` ``Op``, then the output is a graph of the density function for `var` set to the value `rv_value`. Parameters ========== var The `RandomVariable` output that determines the log-likelihood graph. Can also be a list of variables. The final log-likelihood graph will be the sum total of all individual log-likelihood graphs of variables in the list. rv_values A variable, or ``dict`` of variables, that represents the value of `var` in its log-likelihood. If no `rv_value` is provided, ``var.tag.value_var`` will be checked and, when available, used. jacobian Whether or not to include the Jacobian term. scaling A scaling term to apply to the generated log-likelihood graph. transformed Apply transforms. sum Sum the log-likelihood or return each term as a separate list item. """ # TODO: In future when we drop support for tag.value_var most of the following # logic can be removed and logpt can just be a wrapper function that calls aeppl's # joint_logprob directly. # If var is not a list make it one. if not isinstance(var, (list, tuple)): var = [var] # If logpt isn't provided values it is assumed that the tagged value var or # observation is the value variable for that particular RV. if rv_values is None: rv_values = {} for rv in var: value_var = getattr(rv.tag, "observations", getattr(rv.tag, "value_var", None)) if value_var is None: raise ValueError(f"No value variable found for var {rv}") rv_values[rv] = value_var # Else we assume we were given a single rv and respective value elif not isinstance(rv_values, Mapping): if len(var) == 1: rv_values = { var[0]: at.as_tensor_variable(rv_values).astype(var[0].type) } else: raise ValueError( "rv_values must be a dict if more than one var is requested") if scaling: rv_scalings = {} for rv, value_var in rv_values.items(): rv_scalings[value_var] = _get_scaling( getattr(rv.tag, "total_size", None), value_var.shape, value_var.ndim) # Aeppl needs all rv-values pairs, not just that of the requested var. # Hence we iterate through the graph to collect them. tmp_rvs_to_values = rv_values.copy() for node in io_toposort(graph_inputs(var), var): try: curr_vars = [node.default_output()] except ValueError: curr_vars = node.outputs for curr_var in curr_vars: if curr_var in tmp_rvs_to_values: continue # Check if variable has a value variable value_var = getattr(curr_var.tag, "observations", getattr(curr_var.tag, "value_var", None)) if value_var is not None: tmp_rvs_to_values[curr_var] = value_var # After collecting all necessary rvs and values, we check for any value transforms transform_map = {} if transformed: for rv, value_var in tmp_rvs_to_values.items(): if hasattr(value_var.tag, "transform"): transform_map[value_var] = value_var.tag.transform # If the provided value_variable does not have transform information, we # check if the original `rv.tag.value_var` does. # TODO: This logic should be replaced by an explicit dict of # `{value_var: transform}` similar to `rv_values`. else: original_value_var = getattr(rv.tag, "value_var", None) if original_value_var is not None and hasattr( original_value_var.tag, "transform"): transform_map[value_var] = original_value_var.tag.transform transform_opt = TransformValuesOpt(transform_map) temp_logp_var_dict = factorized_joint_logprob(tmp_rvs_to_values, extra_rewrites=transform_opt, use_jacobian=jacobian, **kwargs) # aeppl returns the logpt for every single value term we provided to it. This includes # the extra values we plugged in above, so we filter those we actually wanted in the # same order they were given in. logp_var_dict = {} for value_var in rv_values.values(): logp_var_dict[value_var] = temp_logp_var_dict[value_var] if scaling: for value_var in logp_var_dict.keys(): if value_var in rv_scalings: logp_var_dict[value_var] *= rv_scalings[value_var] if sum: logp_var = at.sum( [at.sum(factor) for factor in logp_var_dict.values()]) else: logp_var = list(logp_var_dict.values()) return logp_var
def logcdfpt( var: TensorVariable, rv_values: Optional[Union[TensorVariable, Dict[TensorVariable, TensorVariable]]] = None, *, scaling: bool = True, sum: bool = True, **kwargs, ) -> TensorVariable: """Create a measure-space (i.e. log-cdf) graph for a random variable at a given point. Parameters ========== var The `RandomVariable` output that determines the log-likelihood graph. rv_values A variable, or ``dict`` of variables, that represents the value of `var` in its log-likelihood. If no `rv_value` is provided, ``var.tag.value_var`` will be checked and, when available, used. jacobian Whether or not to include the Jacobian term. scaling A scaling term to apply to the generated log-likelihood graph. transformed Apply transforms. sum Sum the log-likelihood. """ if not isinstance(rv_values, Mapping): rv_values = {var: rv_values} if rv_values is not None else {} rv_var, rv_value_var = extract_rv_and_value_vars(var) rv_value = rv_values.get(rv_var, rv_value_var) if rv_var is not None and rv_value is None: raise ValueError(f"No value variable specified or associated with {rv_var}") if rv_value is not None: rv_value = at.as_tensor(rv_value) if rv_var is not None: # Make sure that the value is compatible with the random variable rv_value = rv_var.type.filter_variable(rv_value.astype(rv_var.dtype)) if rv_value_var is None: rv_value_var = rv_value rv_node = rv_var.owner rng, size, dtype, *dist_params = rv_node.inputs # Here, we plug the actual random variable into the log-likelihood graph, # because we want a log-likelihood graph that only contains # random variables. This is important, because a random variable's # parameters can contain random variables themselves. # Ultimately, with a graph containing only random variables and # "deterministics", we can simply replace all the random variables with # their value variables and be done. tmp_rv_values = rv_values.copy() tmp_rv_values[rv_var] = rv_var logp_var = _logcdf(rv_node.op, rv_var, tmp_rv_values, *dist_params, **kwargs) transform = getattr(rv_value_var.tag, "transform", None) if rv_value_var else None # Replace random variables with their value variables replacements = rv_values.copy() replacements.update({rv_var: rv_value, rv_value_var: rv_value}) (logp_var,), _ = rvs_to_value_vars( (logp_var,), apply_transforms=False, initial_replacements=replacements, ) if sum: logp_var = at.sum(logp_var) if scaling: logp_var *= _get_scaling( getattr(rv_var.tag, "total_size", None), rv_value.shape, rv_value.ndim ) # Recompute test values for the changes introduced by the replacements # above. if config.compute_test_value != "off": for node in io_toposort(graph_inputs((logp_var,)), (logp_var,)): compute_test_value(node) if rv_var.name is not None: logp_var.name = f"__logp_{rv_var.name}" return logp_var
def on_detach(self, fgraph): for node in io_toposort(fgraph.inputs, fgraph.outputs): self.on_prune(fgraph, node, "Bookkeeper.detach")
def on_attach(self, fgraph): for node in io_toposort(fgraph.inputs, fgraph.outputs): self.on_import(fgraph, node, "on_attach")
def logpt( var: TensorVariable, rv_values: Optional[Union[TensorVariable, Dict[TensorVariable, TensorVariable]]] = None, *, jacobian: bool = True, scaling: bool = True, transformed: bool = True, sum: bool = True, **kwargs, ) -> TensorVariable: """Create a measure-space (i.e. log-likelihood) graph for a random variable or a list of random variables at a given point. The input `var` determines which log-likelihood graph is used and `rv_value` is that graph's input parameter. For example, if `var` is the output of a ``NormalRV`` ``Op``, then the output is a graph of the density function for `var` set to the value `rv_value`. Parameters ========== var The `RandomVariable` output that determines the log-likelihood graph. Can also be a list of variables. The final log-likelihood graph will be the sum total of all individual log-likelihood graphs of variables in the list. rv_values A variable, or ``dict`` of variables, that represents the value of `var` in its log-likelihood. If no `rv_value` is provided, ``var.tag.value_var`` will be checked and, when available, used. jacobian Whether or not to include the Jacobian term. scaling A scaling term to apply to the generated log-likelihood graph. transformed Apply transforms. sum Sum the log-likelihood. """ # TODO: In future when we drop support for tag.value_var most of the following # logic can be removed and logpt can just be a wrapper function that calls aeppl's # joint_logprob directly. # If var is not a list make it one. if not isinstance(var, list): var = [var] # If logpt isn't provided values and the variable (provided in var) # is an RV, it is assumed that the tagged value var or observation is # the value variable for that particular RV. if rv_values is None: rv_values = {} for _var in var: if isinstance(_var.owner.op, RandomVariable): rv_value_var = getattr(_var.tag, "observations", getattr(_var.tag, "value_var", _var)) rv_values = {_var: rv_value_var} elif not isinstance(rv_values, Mapping): # Else if we're given a single value and a single variable we assume a mapping among them. rv_values = ({ var[0]: at.as_tensor_variable(rv_values).astype(var[0].type) } if len(var) == 1 else {}) # Since the filtering of logp graph is based on value variables # provided to this function if not rv_values: warnings.warn( "No value variables provided the logp will be an empty graph") if scaling: rv_scalings = {} for _var in var: rv_value_var = getattr(_var.tag, "observations", getattr(_var.tag, "value_var", _var)) rv_scalings[rv_value_var] = _get_scaling( getattr(_var.tag, "total_size", None), rv_value_var.shape, rv_value_var.ndim) # Aeppl needs all rv-values pairs, not just that of the requested var. # Hence we iterate through the graph to collect them. tmp_rvs_to_values = rv_values.copy() transform_map = {} for node in io_toposort(graph_inputs(var), var): try: curr_vars = [node.default_output()] except ValueError: curr_vars = node.outputs for curr_var in curr_vars: rv_value_var = getattr(curr_var.tag, "observations", getattr(curr_var.tag, "value_var", None)) if rv_value_var is None: continue rv_value = rv_values.get(curr_var, rv_value_var) tmp_rvs_to_values[curr_var] = rv_value # Along with value variables we also check for transforms if any. if hasattr(rv_value_var.tag, "transform") and transformed: transform_map[rv_value] = rv_value_var.tag.transform transform_opt = TransformValuesOpt(transform_map) temp_logp_var_dict = factorized_joint_logprob(tmp_rvs_to_values, extra_rewrites=transform_opt, use_jacobian=jacobian, **kwargs) # aeppl returns the logpt for every single value term we provided to it. This includes # the extra values we plugged in above so we need to filter those out. logp_var_dict = {} for value_var, _logp in temp_logp_var_dict.items(): if value_var in rv_values.values(): logp_var_dict[value_var] = _logp # If it's an empty dictionary the logp is None if not logp_var_dict: logp_var = None else: # Otherwise apply appropriate scalings and at.add and/or at.sum the # graphs accordingly. if scaling: for _value in logp_var_dict.keys(): if _value in rv_scalings: logp_var_dict[_value] *= rv_scalings[_value] if len(logp_var_dict) == 1: logp_var_dict = tuple(logp_var_dict.values())[0] if sum: logp_var = at.sum(logp_var_dict) else: logp_var = logp_var_dict else: if sum: logp_var = at.sum( [at.sum(factor) for factor in logp_var_dict.values()]) else: logp_var = at.add(*logp_var_dict.values()) # Recompute test values for the changes introduced by the replacements # above. if config.compute_test_value != "off": for node in io_toposort(graph_inputs((logp_var, )), (logp_var, )): compute_test_value(node) return logp_var
def logpt( var: TensorVariable, rv_values: Optional[Union[TensorVariable, Dict[TensorVariable, TensorVariable]]] = None, *, jacobian: bool = True, scaling: bool = True, transformed: bool = True, cdf: bool = False, sum: bool = False, **kwargs, ) -> TensorVariable: """Create a measure-space (i.e. log-likelihood) graph for a random variable at a given point. The input `var` determines which log-likelihood graph is used and `rv_value` is that graph's input parameter. For example, if `var` is the output of a ``NormalRV`` ``Op``, then the output is a graph of the density function for `var` set to the value `rv_value`. Parameters ========== var The `RandomVariable` output that determines the log-likelihood graph. rv_values A variable, or ``dict`` of variables, that represents the value of `var` in its log-likelihood. If no `rv_value` is provided, ``var.tag.value_var`` will be checked and, when available, used. jacobian Whether or not to include the Jacobian term. scaling A scaling term to apply to the generated log-likelihood graph. transformed Apply transforms. cdf Return the log cumulative distribution. sum Sum the log-likelihood. """ if not isinstance(rv_values, Mapping): rv_values = {var: rv_values} if rv_values is not None else {} rv_var, rv_value_var = extract_rv_and_value_vars(var) rv_value = rv_values.get(rv_var, rv_value_var) if rv_var is not None and rv_value is None: raise ValueError( f"No value variable specified or associated with {rv_var}") if rv_value is not None: rv_value = at.as_tensor(rv_value) if rv_var is not None: # Make sure that the value is compatible with the random variable rv_value = rv_var.type.filter_variable( rv_value.astype(rv_var.dtype)) if rv_value_var is None: rv_value_var = rv_value if rv_var is None: if var.owner is not None: return _logp( var.owner.op, var, rv_values, *var.owner.inputs, jacobian=jacobian, scaling=scaling, transformed=transformed, cdf=cdf, sum=sum, ) return at.zeros_like(var) rv_node = rv_var.owner rng, size, dtype, *dist_params = rv_node.inputs # Here, we plug the actual random variable into the log-likelihood graph, # because we want a log-likelihood graph that only contains # random variables. This is important, because a random variable's # parameters can contain random variables themselves. # Ultimately, with a graph containing only random variables and # "deterministics", we can simply replace all the random variables with # their value variables and be done. tmp_rv_values = rv_values.copy() tmp_rv_values[rv_var] = rv_var if not cdf: logp_var = _logp(rv_node.op, rv_var, tmp_rv_values, *dist_params, **kwargs) else: logp_var = _logcdf(rv_node.op, rv_var, tmp_rv_values, *dist_params, **kwargs) transform = getattr(rv_value_var.tag, "transform", None) if rv_value_var else None if transform and transformed and not cdf and jacobian: transformed_jacobian = transform.jacobian_det(rv_var, rv_value) if transformed_jacobian: if logp_var.ndim > transformed_jacobian.ndim: logp_var = logp_var.sum(axis=-1) logp_var += transformed_jacobian # Replace random variables with their value variables replacements = rv_values.copy() replacements.update({rv_var: rv_value, rv_value_var: rv_value}) (logp_var, ), _ = rvs_to_value_vars( (logp_var, ), apply_transforms=transformed and not cdf, initial_replacements=replacements, ) if sum: logp_var = at.sum(logp_var) if scaling: logp_var *= _get_scaling(getattr(rv_var.tag, "total_size", None), rv_value.shape, rv_value.ndim) # Recompute test values for the changes introduced by the replacements # above. if config.compute_test_value != "off": for node in io_toposort(graph_inputs((logp_var, )), (logp_var, )): compute_test_value(node) if rv_var.name is not None: logp_var.name = "__logp_%s" % rv_var.name return logp_var