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
