graph_util Module#

concretize_connections(connections, tree, separator='.', in_key='input', out_key='output')[source]#

Concretize the connections dictionary by explicitly adding paths for leaves that have substructures but are not fully specified in the connections. The keys in the connections are never modified, since omission of substructures in the keys implies that the entire substructure is passed. Only the values are modified to explicitly specify all substructures. Order is preserved.

Returns:

  • Concretized connections dictionary with all leave substructures

  • explicitly specified.

Parameters:
Return type:

dict[str, list[str]]

Examples

All examples below assume a tree like
>>> tree = {
...     "M1": "*",
...     "M2": "*",
... }
Multiple implicit connections to the output
>>> connections = {
...     "input": ["M1", "M2"],
...     "M1": "output",
...     "M2": "output"
... }
>>> concretized_connections = {
...     "input": ["M1", "M2"],
...     "M1": ["output.0"],
...     "M2": ["output.1"]
... }
Multiple implicit connections from the input and between modules
>>> connections = {
...     "input.0": "M1",
...     "input.1": "M1",
...     "input.2": "M2",
...     "M1": "M2",
...     "M2": "output"
... }
>>> concretized_connections = {
...     "input.0": ["M1.0"],
...     "input.1": ["M1.1"],
...     "input.2": ["M2.0"],
...     "M1": ["M2.1"],
...     "M2": ["output"]
... }
Beyond depth-1 implicit connections
>>> connections = {
...     "input": ["M1", "M2"],
...     "M1": "output.0",
...     "M2": "output.0"
... }
>>> concretized_connections = {
...     "input": ["M1", "M2"],
...     "M1": ["output.0.0"],
...     "M2": ["output.0.1"]
... }
Partially implicit connection to the output
>>> connections = {
...     "input": ["M1", "M2"],
...     "M1": "output.1",
...     "M2": "output"
... }
>>> concretized_connections = {
...     "input": ["M1", "M2"],
...     "M1": ["output.1"],
...     "M2": ["output.0"]
... }
Partially implicit connections across depths
>>> connections = {
...     "input": ["M1", "M2"],
...     "M1": "output.1.a",
...     "M2": "output"
... }
>>> concretized_connections = {
...     "input": ["M1", "M2"],
...     "M1": ["output.1.a"],
...     "M2": ["output.0"]
... }

Partially implicit connection with a dictionary type, in which case a random UUID is used to avoid collisions

>>> connections = {
...     "input": ["M1", "M2"],
...     "M1": "output.a",
...     "M2": "output"
... }
>>> concretized_connections = {
...     "input": ["M1", "M2"],
...     "M1": ["output.a"],
...     "M2": ["output.<uuid>"]
... }
concretize_paths(paths, separator='.')[source]#

Concretize a list of keypaths by appending indices or uuids to duplicate paths. Respects existing indices or keys.

Parameters:

  • paths (list[str]) – List of keypaths to normalize.

  • separator (str) – The string used to separate levels in the keypaths. Default is “.”.

Returns:

List of concretized keypaths with duplicates resolved.

Return type:

list[str]

Examples

>>> paths = [
...     "a.b.0.1",
...     "a.b.0",
...     "a.b",
...     "a.1",
...     "a.1",
...     "a.2",
...     "a",
...     "",
... ]
>>> concrete_paths = concretize_paths(paths)
>>> concrete_paths
[
    "a.b.0.1",
    "a.b.0.0", # added .0 since a.b.0.* exists
    "a.b.1",   # added .1 since a.b.* exists and .0 is taken
    "a.1.0",   # added .0 since a.1.* exists
    "a.1.1",   # added .1 since a.1.* exists and .0 is taken
    "a.2",     # no change since a.2.* doesn't exist
    "a.<uuid>", # added .<uuid> since a.* exists & non-int keys ('b')
    "<uuid>",  # added <uuid> since * exists & non-int keys ('a')
]
get_outer_connections_by_tree(connections, tree, separator='.', in_key='input', out_key='output')[source]#

Get the connections dictionary showing only the tree structure, not including any substructure.

Returns:

Connections dictionary with only tree structure.

Parameters:
Return type:

dict[str, list[str]]

Examples

Given a tree like
>>> tree = {
...     "block1": {
...         "M1": "*"
...     },
...     "block2": {
...         "M2": "*"
...     },
...     "block3": {
...         "M3": "*"
...     }
... }

and a connections dictionary like

>>> connections = {
...     "block1.M1.a": ["block2.M2.0", "block3.M3.input"],
...     "block1.M1.b": "output",
...     "input": "block1.M1"
... }

The outer tree-only connections will be

>>> outer_connections = {
...     "block1.M1": ["block2.M2", "block3.M3"],
...     "block1.M1": ["output"],
...     "input": ["block1.M1"]
... }
partition_connections_by_tree(connections, tree, separator='.', in_key='input', out_key='output')[source]#

Process the connections dictionary to separate the tree structure of ‘tree’ from the remaining structure.

Parameters:

  • connections (dict[str, list[str]]) – Dictionary defining connections between leaves in the tree and their sub-structures. Keys and values are strings representing the keypaths to the leaves or sub-structures in the tree.

  • tree (PyTree[Any]) – A PyTree representing the outer structure to separate.

  • separator (str) – The string used to separate levels in the keypaths. Default is “.”.

  • in_key (str) – The reserved key representing the model input. Default is “input”.

  • out_key (str) – The reserved key representing the model output. Default is “output”.

Returns:

  • Processed connections dictionary with separated structures

  • by a double separator (e.g., “..” if the separator is “.”).

Raises:

ValueError – If the connections contain invalid keys or values, or if the separator appears consecutively in any key or value.

Return type:

dict[str, list[str]]

Examples

Given a tree like
>>> tree = {
...     "block1": {
...         "M1": "*"
...     },
...     "block2": {
...         "M2": "*"
...     },
...     "block3": {
...         "M3": "*"
...     }
... }
and connections dictionary like
>>> connections = {
...     "block1.M1.a": ["block2.M2.0", "block3.M3.input"],
...     "block1.M1.b": "output",
...     "input": "block1.M1"
... }
The partitioned connections will be
>>> partitioned_connections = {
...     "block1..M1.a": ["block2..M2.0", "block3..M3.input"],
...     "block1..M1.b": ["output.."],
...     "input..": ["block1..M1"]
... }
place_connections_in_tree(connections, tree, separator='.', in_key='input', out_key='output')[source]#

Place the concretized connections into a PyTree structure matching ‘tree’. :param connections: Connections dictionary, will be concretized internally if needed. :param tree: A PyTree representing the outer structure to separate. :param separator: The string used to separate levels in the keypaths. Default is “.”.

Returns:

  • A PyTree with the same structure as ‘tree’, with values from

  • ’concrete_connections’ placed at the appropriate leaves. And the

  • remainder structure with the end_key paths.

Parameters:
Return type:

tuple[PyTree[Any], PyTree[Any]]

Examples

Given a tree like
>>> tree = {
...     "M1": "*",
...     "M2": "*",
... }
and concretized connections like
>>> concretized_connections = {
...     "input": ["M1", "M2"],
...     "M1": ["output.0"],
...     "M2": ["output.1"]
... }
this will (internally) reverse the connections to get
>>> reversed_connections = {
...    "M1": ["input"],
...    "M2": ["input"],
...    "output.0": ["M1"],
...    "output.1": ["M2"],
... }

then place the paths into the original tree structure, and the remainder with the end_key will be returned separately, with its corresponding structure

>>> placed_tree = {
...     "M1": ["input"],
...     "M2": ["input"],
... }
>>> output_remainder = [
...     "M1",
...     "M2",
... ]
Further examples:
>>> tree = {"A": "*", "B": ("*", "*")}
>>> concretized_connections = {
...     "input.0": ["A.0", "B.0", "B.1"],
...     "input.1": ["A.1"],
...     "B.0": ["A.2"],
...     "B.1": ["output.x"],
...     "A": ["output.y"],
... }
>>> placed_tree, output_remainder = place_connections_in_tree(
...     concretized_connections,
...     tree,
... )
>>> placed_tree
{
    "A": ["input.0", "input.1", "B.0"],
    "B": ("input.0", "input.0"),
}
>>> output_remainder
{
    "x": "B.1",
    "y": "A",
}
resolve_connections(graph, start_key='input', end_key='output')[source]#
Parameters:
Return type:

tuple[list[str], OrderedSet[str]]

validate_resolution(order, all_nodes, start_key='input', end_key='output')[source]#
Parameters:
Return type:

None