def copy_dict_of_named_arrays(source_dict: DictOfNamedArrays,
copy_mapper: CopyMapper) -> DictOfNamedArrays:
"""Copy the elements of a :class:`~pytato.DictOfNamedArrays` into a
:class:`~pytato.DictOfNamedArrays` with a new namespace.
:param source_dict: The :class:`~pytato.DictOfNamedArrays` to copy
:param copy_mapper: A mapper that performs copies into a new namespace
:returns: A new :class:`~pytato.DictOfNamedArrays` containing copies of the
items in *source_dict*
"""
from pytato.scalar_expr import get_dependencies as get_scalar_expr_deps
if not source_dict:
return DictOfNamedArrays({})
# {{{ extract dependencies elements of the namespace
# https://github.com/python/mypy/issues/2013
deps: FrozenSet[Array] = frozenset().union(*list(
get_dependencies(source_dict).values())) # type: ignore
dep_names = (frozenset(
[dep.name for dep in deps if isinstance(dep, InputArgumentBase)])
| get_scalar_expr_deps([dep.shape for dep in deps]))
# }}}
data = {}
copy_namespace(source_dict.namespace, copy_mapper, dep_names)
for name, val in source_dict.items():
data[name] = copy_mapper(val)
return DictOfNamedArrays(data)
def copy_dict_of_named_arrays(source_dict: DictOfNamedArrays,
copy_mapper: CopyMapper) -> DictOfNamedArrays:
"""Copy the elements of a :class:`~pytato.DictOfNamedArrays` into a
:class:`~pytato.DictOfNamedArrays` with a new namespace.
:param source_dict: The :class:`~pytato.DictOfNamedArrays` to copy
:param copy_mapper: A mapper that performs copies into a new namespace
:returns: A new :class:`~pytato.DictOfNamedArrays` containing copies of the
items in *source_dict*
"""
if not source_dict:
return DictOfNamedArrays({})
data = {}
copy_namespace(source_dict.namespace, copy_mapper)
for name, val in source_dict.items():
data[name] = copy_mapper(val)
return DictOfNamedArrays(data)
def normalize_outputs(result: Union[Array, DictOfNamedArrays,
Dict[str, Array]]) -> DictOfNamedArrays:
"""Convert outputs of a computation to the canonical form.
Performs a conversion to :class:`~pytato.DictOfNamedArrays` if necessary.
:param result: Outputs of the computation.
"""
if not isinstance(result, (Array, DictOfNamedArrays, dict)):
raise TypeError("outputs of the computation should be "
"either an Array or a DictOfNamedArrays")
if isinstance(result, Array):
outputs = DictOfNamedArrays({"_pt_out": result})
elif isinstance(result, dict):
outputs = DictOfNamedArrays(result)
else:
assert isinstance(result, DictOfNamedArrays)
outputs = result
return outputs
def preprocess(outputs: DictOfNamedArrays) -> PreprocessResult:
"""Preprocess a computation for code generation."""
from pytato.transform import copy_dict_of_named_arrays, get_dependencies
# {{{ compute the order in which the outputs must be computed
# semantically order does not matter, but doing a toposort ordering of the
# outputs leads to a FLOP optimal choice
from pytools.graph import compute_topological_order
deps = get_dependencies(outputs)
# only look for dependencies between the outputs
deps = {
name: (val & frozenset(outputs.values()))
for name, val in deps.items()
}
# represent deps in terms of output names
output_to_name = {output: name for name, output in outputs.items()}
dag = {
name: (frozenset([output_to_name[output]
for output in val]) - frozenset([name]))
for name, val in deps.items()
}
output_order: List[str] = compute_topological_order(dag)[::-1]
# }}}
mapper = CodeGenPreprocessor(Namespace())
new_outputs = copy_dict_of_named_arrays(outputs, mapper)
return PreprocessResult(outputs=new_outputs,
compute_order=tuple(output_order),
bound_arguments=mapper.bound_arguments)
def generate_code_for_partition(partition: GraphPartition) \
-> Dict[PartId, BoundProgram]:
"""Return a mapping of partition identifiers to their
:class:`pytato.target.BoundProgram`."""
from pytato import generate_loopy
part_id_to_prg = {}
for part in partition.parts.values():
d = DictOfNamedArrays(
{var_name: partition.var_name_to_result[var_name]
for var_name in part.output_names
})
part_id_to_prg[part.pid] = generate_loopy(d)
return part_id_to_prg
def check_validity_of_outputs(exprs: DictOfNamedArrays) -> None:
name_validation_mapper = NamesValidityChecker()
for ary in exprs.values():
name_validation_mapper(ary)
def get_dependencies(expr: DictOfNamedArrays) -> Dict[str, FrozenSet[Array]]:
"""Returns the dependencies of each named array in *expr*.
"""
dep_mapper = DependencyMapper()
return {name: dep_mapper(val) for name, val in expr.items()}