Transforming Array Expression Graphs

class pytato.transform.Mapper

A class that when called with a pytato.Array recursively iterates over the DAG, calling the _mapper_method of each node. Users of this class are expected to override the methods of this class or create a subclass.

Note

This class might visit a node multiple times. Use a CachedMapper if this is not desired.

handle_unsupported_array(expr: pytato.transform.T, *args: Any, **kwargs: Any) → Any

Mapper method that is invoked for pytato.Array subclasses for which a mapper method does not exist in this mapper.

map_foreign(expr: Any, *args: Any, **kwargs: Any) → Any

Mapper method that is invoked for an object of class for which a mapper method does not exist in this mapper.

rec(expr: pytato.transform.T, *args: Any, **kwargs: Any) → Any

Call the mapper method of expr and return the result.

__call__(expr: pytato.transform.T, *args: Any, **kwargs: Any) → Any

Handle the mapping of expr.

class pytato.transform.CachedMapper

Mapper class that maps each node in the DAG exactly once. This loses some information compared to Mapper as a node is visited only from one of its predecessors.

class pytato.transform.CopyMapper

Performs a deep copy of a pytato.array.Array. The typical use of this mapper is to override individual map_ methods in subclasses to permit term rewriting on an expression graph.

Note

This does not copy the data of a pytato.array.DataWrapper.

class pytato.transform.CombineMapper

class pytato.transform.DependencyMapper

Maps a pytato.array.Array to a frozenset of pytato.array.Array’s it depends on.

Warning

This returns every node in the graph! Consider a custom CombineMapper or a SubsetDependencyMapper instead.

class pytato.transform.InputGatherer

Mapper to combine all instances of pytato.array.InputArgumentBase that an array expression depends on.

class pytato.transform.SizeParamGatherer

Mapper to combine all instances of pytato.array.SizeParam that an array expression depends on.

class pytato.transform.SubsetDependencyMapper(universe: FrozenSet[pytato.array.Array])

Mapper to combine the dependencies of an expression that are a subset of universe.

class pytato.transform.WalkMapper

A mapper that walks over all the arrays in a pytato.Array.

Users may override the specific mapper methods in a derived class or override WalkMapper.visit() and WalkMapper.post_visit().

visit(expr: Any) → bool

If this method returns True, expr is traversed during the walk. If this method returns False, expr is not traversed as a part of the walk.

post_visit(expr: Any) → None

Callback after expr has been traversed.

class pytato.transform.CachedWalkMapper

WalkMapper that visits each node in the DAG exactly once. This loses some information compared to WalkMapper as a node is visited only from one of its predecessors.

class pytato.transform.TopoSortMapper

A mapper that creates a list of nodes in topological order.

Members

topological_order

class pytato.transform.CachedMapAndCopyMapper(map_fn: Callable[[Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays]], Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays]])

Mapper that applies map_fn to each node and copies it. Results of traversals are memoized i.e. each node is mapped via map_fn exactly once.

class pytato.transform.EdgeCachedMapper

Mapper class to execute a rewriting method (handle_edge()) on each edge in the graph.

abstract handle_edge(expr: Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays], child: Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays]) → Any

pytato.transform.copy_dict_of_named_arrays(source_dict: pytato.array.DictOfNamedArrays, copy_mapper: pytato.transform.CopyMapper) → pytato.array.DictOfNamedArrays

Copy the elements of a DictOfNamedArrays into a DictOfNamedArrays.

Parameters

• source_dict – The DictOfNamedArrays to copy

• copy_mapper – A mapper that performs copies different array types

Returns

A new DictOfNamedArrays containing copies of the items in source_dict

pytato.transform.get_dependencies(expr: pytato.array.DictOfNamedArrays) → Dict[str, FrozenSet[pytato.array.Array]]

Returns the dependencies of each named array in expr.

pytato.transform.map_and_copy(expr: pytato.transform.T, map_fn: Callable[[Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays]], Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays]]) → Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays]

Returns a copy of expr with every array expression reachable from expr mapped via map_fn.

Note

Uses CachedMapAndCopyMapper under the hood and because of its caching nature each node is mapped exactly once.

pytato.transform.materialize_with_mpms(expr: pytato.array.DictOfNamedArrays) → pytato.array.DictOfNamedArrays

Materialize nodes in expr with MPMS materialization strategy. MPMS stands for Multiple-Predecessors, Multiple-Successors.

Note

• MPMS materialization strategy is a greedy materialization algorithm in which any node with more than 1 materialized predecessors and more than 1 successors is materialized.

• Materializing here corresponds to tagging a node with ImplStored.

• Does not attempt to materialize sub-expressions in pytato.Array.shape.

Warning

This is a greedy materialization algorithm and thereby this algorithm might be too eager to materialize. Consider the graph below:

I1          I2
 \         /
  \       /
   \     /
    🡦   🡧
      T
     / \
    /   \
   /     \
  🡧       🡦
 O1        O2

where, ‘I1’, ‘I2’ correspond to instances of pytato.array.InputArgumentBase, and, ‘O1’ and ‘O2’ are the outputs required to be evaluated in the computation graph. MPMS materialization algorithm will materialize the intermediate node ‘T’ as it has 2 predecessors and 2 successors. However, the total number of memory accesses after applying MPMS goes up as shown by the table below.

	Before	After
Reads	4	4
Writes	2	3
Total	6	7

Dict representation of DAGs

class pytato.transform.UsersCollector

Maps a graph to a dictionary representation mapping a node to its users, i.e. all the nodes using its value.

node_to_users

Mapping of each node in the graph to its users.

__init__() → None

pytato.transform.reverse_graph(graph: Dict[Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays], Set[Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays]]]) → Dict[Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays], Set[Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays]]]

Reverses a graph.

Parameters

graph – A dict representation of a directed graph, mapping each node to other nodes to which it is connected by edges. A possible use case for this function is the graph in UsersCollector.node_to_users.

Returns

A dict representing graph with edges reversed.

pytato.transform.tag_child_nodes(graph: Dict[Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays], Set[Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays]]], tag: Any, starting_point: Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays], node_to_tags: Optional[Dict[Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays], Set[Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays]]]] = None) → Dict[Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays], Set[Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays]]]

Tags all nodes reachable from starting_point with tag.

Parameters

• graph – A dict representation of a directed graph, mapping each node to other nodes to which it is connected by edges. A possible use case for this function is the graph in UsersCollector.node_to_users.

• tag – The value to tag the nodes with.

• starting_point – An optional starting point in graph.

• node_to_tags – The resulting mapping of nodes to tags.

Returns

the updated value of node_to_tags.

Internal stuff that is only here because the documentation tool wants it

class pytato.transform.T

A type variable representing the input type of a Mapper.

Analyzing Array Expression Graphs

pytato.analysis.get_nusers(outputs: Union[pytato.array.Array, pytato.array.DictOfNamedArrays]) → Mapping[pytato.array.Array, int]

For the DAG outputs, returns the mapping from each node to the number of nodes using its value within the DAG given by outputs.

pytato.analysis.is_einsum_similar_to_subscript(expr: pytato.array.Einsum, subscripts: str) → bool

Returns True if and only if an einsum with the subscript descriptor string subscripts operated on expr’s pytato.array.Einsum.args would compute the same result as expr.

pytato.analysis.get_num_nodes(outputs: Union[pytato.array.Array, pytato.array.DictOfNamedArrays]) → int

Returns the number of nodes in DAG outputs.

Visualizing Array Expression Graphs

pytato.get_dot_graph(result: Union[pytato.array.Array, pytato.array.DictOfNamedArrays]) → str

Return a string in the dot language depicting the graph of the computation of result.

Parameters

result – Outputs of the computation (cf. pytato.generate_loopy()).

pytato.get_dot_graph_from_partition(partition: pytato.partition.GraphPartition) → str

Return a string in the dot language depicting the graph of the partitioned computation of partition.

Parameters

partition – Outputs of find_partition().

pytato.show_dot_graph(result: Union[str, pytato.array.Array, pytato.array.DictOfNamedArrays, pytato.partition.GraphPartition]) → None

Show a graph representing the computation of result in a browser.

Parameters

result – Outputs of the computation (cf. pytato.generate_loopy()) or the output of get_dot_graph(), or the output of find_partition().

pytato.get_ascii_graph(result: Union[pytato.array.Array, pytato.array.DictOfNamedArrays], use_color: bool = True) → str

Return a string representing the computation of result using the asciidag package.

Parameters

• result – Outputs of the computation (cf. pytato.generate_loopy()).

• use_color – Colorized output

pytato.show_ascii_graph(result: Union[pytato.array.Array, pytato.array.DictOfNamedArrays]) → None

Show a graph representing the computation of result in a browser.

Parameters

result – Outputs of the computation (cf. pytato.generate_loopy()) or the output of get_dot_graph().

Comparing two expression Graphs

class pytato.equality.EqualityComparer

A pytato.array.Array visitor to check equality between two expression DAGs.

Note

• Compares two expression graphs expr1, expr2 in \(O(N)\) comparisons, where \(N\) is the number of nodes in expr1.

• This visitor was introduced to memoize the sub-expression comparisons of the expressions to be compared. Not memoizing the sub-expression comparisons results in \(O(2^N)\) complexity for the comparison operation, where \(N\) is the number of nodes in expressions. See GH-Issue-163 <https://github.com/inducer/pytato/issues/163> for more on this.

Stringifying Expression Graphs

class pytato.stringifier.Reprifier(truncation_depth: int = 3, truncation_string: str = '(...)')

Stringifies pytato-types to closely resemble CPython’s implementation of repr() for its builtin datatypes.

Partitioning Array Expression Graphs

class pytato.partition.GraphPart(pid: Hashable, needed_pids: FrozenSet[Hashable], input_names: FrozenSet[str], output_names: FrozenSet[str])

pid

An identifier for this part of the graph.

needed_pids

The IDs of parts that are required to be evaluated before this part can be evaluated.

input_names

Names of placeholders the part requires as input.

output_names

Names of placeholders this part provides as output.

class pytato.partition.GraphPartition(parts: Mapping[Hashable, pytato.partition.GraphPart], var_name_to_result: Mapping[str, pytato.array.Array], toposorted_part_ids: List[Hashable])

Store information about a partitioning of an expression graph.

parts

Mapping from part IDs to instances of GraphPart.

var_name_to_result

Mapping of placeholder names to the respective pytato.array.Array they represent.

toposorted_part_ids

One possible topologically sorted ordering of part IDs that is admissible under GraphPart.needed_pids.

Note

This attribute could be recomputed for those dependencies. Since it is computed as part of find_partition() anyway, it is preserved here.

exception pytato.partition.PartitionInducedCycleError

Raised by find_partition() if the partitioning induced a cycle in the graph of partitions.

pytato.partition.find_partition(outputs: pytato.array.DictOfNamedArrays, part_func: Callable[[Union[pytato.array.Array, pytato.array.AbstractResultWithNamedArrays]], Hashable]) → pytato.partition.GraphPartition

Partitions the expr according to part_func and generates code for each partition. Raises PartitionInducedCycleError if the partitioning induces a cycle, e.g. for a graph like the following:

   ┌───┐
┌──┤ A ├──┐
│  └───┘  │
│       ┌─▼─┐
│       │ B │
│       └─┬─┘
│  ┌───┐  │
└─►│ C │◄─┘
   └───┘

where A and C are in partition 1, and B is in partition 2.

Parameters

• expr – The expression to partition.

• part_func – A callable that returns an instance of Hashable for a node.

Returns

An instance of GraphPartition that contains the partition.

pytato.partition.execute_partition(partition: pytato.partition.GraphPartition, prg_per_partition: Dict[Hashable, pytato.target.BoundProgram], queue: Any) → Dict[str, Any]

Executes a set of partitions on a pyopencl.CommandQueue.

Parameters

• parts – An instance of GraphPartition representing the partitioned code.

• queue – An instance of pyopencl.CommandQueue to execute the code on.

Returns

A dictionary of variable names mapped to their values.

Utilities and Diagnostics

Helper routines

pytato.utils.are_shape_components_equal(dim1: Union[int, numpy.int8, numpy.int16, numpy.int32, numpy.int64, numpy.uint8, numpy.uint16, numpy.uint32, numpy.uint64, pytato.array.Array], dim2: Union[int, numpy.int8, numpy.int16, numpy.int32, numpy.int64, numpy.uint8, numpy.uint16, numpy.uint32, numpy.uint64, pytato.array.Array]) → bool

Returns True iff dim1 and dim2 are have equal SizeParam coefficients in their expressions.

pytato.utils.are_shapes_equal(shape1: Tuple[Union[int, numpy.int8, numpy.int16, numpy.int32, numpy.int64, numpy.uint8, numpy.uint16, numpy.uint32, numpy.uint64, pytato.array.Array], ...], shape2: Tuple[Union[int, numpy.int8, numpy.int16, numpy.int32, numpy.int64, numpy.uint8, numpy.uint16, numpy.uint32, numpy.uint64, pytato.array.Array], ...]) → bool

Returns True iff shape1 and shape2 have the same dimensionality and the correpsonding components are equal as defined by are_shape_components_equal().

pytato.utils.get_shape_after_broadcasting(exprs: Sequence[Union[pytato.array.Array, numbers.Number, int, numpy.bool_, bool, pymbolic.primitives.Expression]]) → Tuple[Union[int, numpy.int8, numpy.int16, numpy.int32, numpy.int64, numpy.uint8, numpy.uint16, numpy.uint32, numpy.uint64, pytato.array.Array], ...]

Returns the shape after broadcasting exprs in an operation.

pytato.utils.dim_to_index_lambda_components(expr: Union[int, numpy.int8, numpy.int16, numpy.int32, numpy.int64, numpy.uint8, numpy.uint16, numpy.uint32, numpy.uint64, pytato.array.Array], vng: Optional[pytools.UniqueNameGenerator] = None) → Tuple[Union[numbers.Number, int, numpy.bool_, bool, pymbolic.primitives.Expression], Dict[str, pytato.array.SizeParam]]

Returns the scalar expressions and bindings to use the shape component within an index lambda.

>>> n = pt.make_size_param("n")
>>> expr, bnds = dim_to_index_lambda_components(3*n+8, UniqueNameGenerator())
>>> print(expr)
3*_in + 8
>>> bnds
{'_in': SizeParam(name='n')}

Pytato-specific exceptions

class pytato.diagnostic.NameClashError

Raised when 2 non-identical InputArgumentBase’s are reachable in an Array’s DAG and share the same name. Here, we refer to 2 objects a and b as being identical iff a is b.

class pytato.diagnostic.CannotBroadcastError