import copy
from opto.trace.containers import NodeContainer
from opto.trace.nodes import Node
[docs]
def recursive_conversion(true_func, false_func):
"""Recursively apply true_func to the nodes and false_func to the rest of
the objects in a container of nodes. Container of nodes are tuple, list,
dict, set, and NodeContainer.
Args:
true_func (callable): the function to be applied to the nodes.
false_func (callable): the function to be applied to the rest of the objects.
"""
def func(obj):
if isinstance(obj, Node): # base case
return true_func(obj)
elif isinstance(obj, tuple):
return tuple(func(x) for x in obj)
elif isinstance(obj, list):
return [func(x) for x in obj]
elif isinstance(obj, dict):
return {k: func(v) for k, v in obj.items()}
elif isinstance(obj, set):
return {func(x) for x in obj}
elif isinstance(obj, NodeContainer):
output = copy.copy(obj)
for k, v in obj.__dict__.items():
setattr(output, k, func(v))
return output
else:
return false_func(obj)
return func
# TODO to test it and clean up the code
[docs]
def apply_op(op, output, *args, **kwargs):
"""A broadcasting operation that applies an op to container of Nodes.
Args:
op (callable): the operator to be applied.
output (Any): the container to be updated.
*args (Any): the positional inputs of the operator.
**kwargs (Any): the keyword inputs of the operator.
"""
inputs = list(args) + list(kwargs.values())
containers = [x for x in inputs if not isinstance(x, Node)]
if len(containers) == 0: # all inputs are Nodes, we just apply op
return op(*args, **kwargs)
# # there is at least one container
# output = copy.deepcopy(containers[0]) # this would be used as the template of the output
def admissible_type(x, base):
return type(x) is type(base) or isinstance(x, Node)
assert all(
admissible_type(x, output) for x in inputs
) # All inputs are either Nodes or the same type as output
if isinstance(output, list) or isinstance(output, tuple):
assert all(
isinstance(x, Node) or len(output) == len(x) for x in inputs
), f"output {output} and inputs {inputs} are of different lengths."
for k in range(len(output)):
_args = [x if isinstance(x, Node) else x[k] for x in args]
_kwargs = {
kk: vv if isinstance(vv, Node) else vv[k] for kk, vv in kwargs.items()
}
output[k] = apply_op(op, output[k], *_args, **_kwargs)
if isinstance(output, tuple):
output = tuple(output)
elif isinstance(output, dict):
for k, v in output.items():
_args = [x if isinstance(x, Node) else x[k] for x in args]
_kwargs = {
kk: vv if isinstance(vv, Node) else vv[k] for kk, vv in kwargs.items()
}
output[k] = apply_op(op, output[k], *_args, **_kwargs)
elif isinstance(output, NodeContainer): # this is a NodeContainer object instance
for k, v in output.__dict__.items():
_args = [x if isinstance(x, Node) else getattr(x, k) for x in args]
_kwargs = {
kk: vv if isinstance(v, Node) else getattr(vv, k)
for kk, vv in kwargs.items()
}
new_v = apply_op(op, v, *_args, **_kwargs)
setattr(output, k, new_v)
else:
pass
return output
