opto.trace.nodes

Module Contents

Classes

Graph	Graph is a registry of all the nodes, forming a Directed Acyclic Graph (DAG).
AbstractNode	AbstractNode represents an abstract data node in a directed graph.
NodeVizStyleGuide	A class to provide a standardized way to visualize nodes in a graph, particularly for use with graph visualization tools like Graphviz.
NodeVizStyleGuideColorful	A class to provide a colorful style guide for visualizing nodes in a graph.
Node	A data node in a directed graph, this is a basic data structure of Trace.
ParameterNode
MessageNode	Output of an operator.
ExceptionNode	Node containing the exception message.

Functions

node	Create a Node object from data. If the data is already a Node, it will be returned as is. This function is provided for the convenience of the user and should be used instead of directly invoking the Node class.
get_op_name	Extract the operator type from the description.

Data

NAME_SCOPES
GRAPH
USED_NODES
T	Graph is a registry of all the nodes, forming a Directed Acyclic Graph (DAG).
IDENTITY_OPERATORS

API

node(data, name=None, trainable=False, description=None, constraint=None)

Create a Node object from data. If the data is already a Node, it will be returned as is. This function is provided for the convenience of the user and should be used instead of directly invoking the Node class.

data: The data to create the Node from.

name: (optional) The name of the Node.

trainable: (optional) A boolean indicating whether the Node is trainable or not. Default is False.

description: (optional) A string describing the data.

constraint: (optional) A string describing any constraint that the data should obey.

The node function allows users to create Node objects from data. The function first checks if the trainable parameter is True. If it is, it checks if the data is already a Node. If it is, it extracts the underlying data and updates the name if a new name is provided. It then creates a ParameterNode object with the extracted data, name, trainable set to True, and the provided constraint. If the message is not already a Node, it creates a new ParameterNode object with the message as the data, the provided name, trainable set to True, and the provided constraint.

If the trainable parameter is False, the function checks if the message is already a Node. If it is, it checks if a name is provided. If a name is provided, it issues a warning that the name is ignored because the message is already a Node. It then returns the message as is. If the message is not already a Node, it creates a new Node object with the message as the data, the provided name, and the provided constraint.

NAME_SCOPES = []

class Graph

Graph is a registry of all the nodes, forming a Directed Acyclic Graph (DAG).

TRACE: A class-level boolean attribute that determines whether the graph is traced when creating MessageNode. Default is True.

_nodes: An instance-level attribute, which is a defaultdict of lists, used as a lookup table to find nodes by name.

The Graph class manages and organizes nodes in a Directed Acyclic Graph (DAG). It provides methods to register nodes, clear the graph, retrieve nodes by name, and identify root nodes.

The register method assumes that elements in _nodes are never removed, which is important for maintaining the integrity of node names.

Initialization

Initialize the Graph object, setting up the _nodes attribute as a defaultdict of lists to store nodes by their names.

TRACE = True

clear()

Remove all nodes from the graph by deleting each node and reinitializing the _nodes attribute.

The clear function iterates over the current nodes stored in the _nodes attribute and deletes each node. After all nodes have been deleted, it reinitializes the _nodes attribute to an empty defaultdict of lists. This ensures that the graph is completely cleared and ready to be repopulated with new nodes if necessary.

The function is called in unit tests to reset the state of the graph between test cases, ensuring that each test runs with a clean slate and is not affected by the state left by previous tests.

After calling clear, any references to the previously stored nodes will become invalid.

register(node)

Add a node to the graph, ensuring that the node is an instance of the Node class and that its name follows the expected format (containing a colon). This method also handles name scoping and assigns a unique name to the node based on its position in the DAG.

node: The node object to be registered in the graph.

After checking that the input is a Node and its name has the right format, the function splits the name of the node into the name variable and the identifier. The function then checks if there are any name scopes defined in the NAME_SCOPES list. If the length of the list is greater than 0, the name is prefixed with the last scope in the list followed by a “/”. This allows for scoping of node names. Finally, the function adds the node to the _nodes dictionary using the modified name as the key. The _name attribute of the node is set to the modified name followed by the index of the node in the list of nodes with the same name.

The register function should only be called after the node has been properly initialized and its name has been set. The function assumes that elements in the _nodes dictionary never get removed.

get(name)

The get method retrieves a node from the graph by its name, which includes an identifier.

name: A string in the format “name:id”, where “name” is the name of the node and “id” is the identifier of the node.

The get function is designed to extract and return a specific node from the graph. The input parameter ‘name’ is expected to be a string formatted as “name:id”. The function first splits this string into two parts: ‘name’ and ‘id’, using the colon (“:”) as the delimiter. The ‘name’ part represents the name of the node, and the ‘id’ part represents the identifier of the node, which is then converted to an integer. The function then accesses the ‘_nodes’ dictionary attribute of the graph object, using the ‘name’ as the key to retrieve the list of nodes associated with that name. Finally, it returns the node at the position specified by the integer ‘id’ within that list.

Ensure that the ‘name’ parameter is correctly formatted as “name:id” before calling this function. The function assumes that the ‘_nodes’ attribute is a dictionary where each key is a node name and the corresponding value is a list of nodes. The ‘id’ should be a valid index within the list of nodes for the given ‘name’.

property roots

The roots property returns a list of all root nodes in the graph. A root node is identified by its is_root attribute.

GRAPH = 'Graph(...)'

USED_NODES = 'list(...)'

T = 'TypeVar(...)'

Graph is a registry of all the nodes, forming a Directed Acyclic Graph (DAG).

TRACE: A class-level boolean attribute that determines whether the graph is traced when creating MessageNode. Default is True.

_nodes: An instance-level attribute, which is a defaultdict of lists, used as a lookup table to find nodes by name.

The Graph class manages and organizes nodes in a Directed Acyclic Graph (DAG). It provides methods to register nodes, clear the graph, retrieve nodes by name, and identify root nodes.

The register method assumes that elements in _nodes are never removed, which is important for maintaining the integrity of node names.

class AbstractNode(value, *, name=None, trainable=False)

Bases: typing.Generic[opto.trace.nodes.T]

AbstractNode represents an abstract data node in a directed graph.

data: The data stored in the node. parents: The list of parent nodes. children: The list of child nodes. name: The name of the node. py_name: The name of the node without the “:” character. id: The ID of the node. level: The level of the node in the graph. is_root: A boolean indicating whether the node is a root node. is_leaf: A boolean indicating whether the node is a leaf node.

The AbstractNode class is meant to be subclassed and extended to create specific types of nodes. The node can have multiple parents and children, forming a directed graph structure. The node has a name, which is used to identify it within the graph. The py_name attribute is the same as the name attribute, but with the “:” character removed.

The node can be initialized with a value, an optional name, and an optional trainable flag. If the value is an instance of the Node class, the node will be initialized as a reference to that node, otherwise, the value will be stored directly in the node. The default name is generated based on the type of the value and a version number which serves as the identifier, separated by “:”.

The AbstractNode class provides several properties to access its attributes. The data property allows access to the stored data. If the node is being traced within a context, the data property adds the node to the list of nodes used in that context. The parents property returns a list of parent nodes, and the children property returns a list of child nodes. The name property returns the name of the node, and the py_name property returns the name without the “:” character. The id property returns the version number/identifier extracted from the name. The level property returns the level of the node in the DAG. The is_root property returns True if the node has no parents, and the is_leaf property returns True if the node has no children.

The AbstractNode class also provides internal methods to add parents and children to the node. The _add_child method adds a child node to the node’s list of children. The _add_parent method adds a parent node to the node’s list of parents and updates the level of the node based on the parent’s level.

The AbstractNode class overrides the __str__ method to provide a string representation of the node. The representation includes the name, the type of the data, and the data itself. The AbstractNode class implements the __deepcopy__ method to create a deep copy of the node. This allows the node to be detached from the original graph. The AbstractNode class provides comparison methods lt and gt to compare the levels of two nodes in the DAG.

Initialization

Initialize an instance of the AbstractNode class.

value: The value to be assigned to the node. name: The name of the node (optional). trainable: A boolean indicating whether the node is trainable or not (optional).

During initialization, this function generates a default name for the node based on the type of the value parameter. If the name parameter is provided, it is appended to the default name. The format of the name is “type:version”, where the version is set to 0 if no name is provided. If the value parameter is an instance of the Node class, the _data attribute of the current node is set to the _data attribute of the value parameter, and the _name attribute is set to the _name attribute of the value parameter if no name is provided. Otherwise, the _data attribute is set to the value parameter itself, and the _name attribute is set to the default name. Finally, the function calls the register function of the GRAPH object to register the current node in the graph.

property data

Retrieve the internal data of a node, potentially adding the node to a list of used nodes if certain conditions are met.

This function assumes that the “_data” attribute exists within the node object. If this attribute is not present, an AttributeError will be raised.

property parents

Access the parents of a node. It is an essential part of the graph structure and is used in various operations such as graph traversal and feedback propagation.

property children

Access the children of a node. This property is essential for accessing the hierarchical structure of nodes, allowing traversal and manipulation of the DAG.

property name

This property is set when the node is registered in the graph. It is a combination of the node’s name and its index in the list of nodes with the same name. The index is incremented each time a new node with the same name is registered. This assumes that elements in the _nodes dictionary of the graph never get removed.

property py_name

property id

The name property is a string formatted as “name:identifier”. This property splits that string using the colon (“:”) delimiter and returns the second part, which corresponds to the identifier. This identifier is typically a unique part of the node’s name, distinguishing it from other nodes with the same base name. Ensure that the name attribute contains a colon (“:”) to avoid index errors during the split operation.

property level

The level of a node in the graph. The level is determined by the maximum level of its parents plus one. The level of a root node is 0.

property is_root

A boolean indicating whether the node is a root node in a graph structure. A root node has no parents.

property is_leaf

A boolean indicating whether the node is a leaf node in a graph structure. A leaf node has no children.

lt(other)

Less than comparison based on the level attribute of the nodes.

other: The other node to compare against.

This method is used to compare the levels of two nodes in the DAG. Therefore it checks if the negated level of the current node (-self._level) is less than the negated level of the other node (-other._level)

gt(other)

Greater than comparison based on the level attribute of the nodes.

other: The other node to compare against.

This method is used to compare the levels of two nodes in the DAG. Therefore it checks if the negated level of the current node (-self._level) is greater than the negated level of the other node (-other._level)

IDENTITY_OPERATORS = ('identity', 'clone')

get_op_name(description)

Extract the operator type from the description.

description: A string containing the description of the node.

The get_op_name function takes a description as input and uses regular expression to search for the operator type enclosed in square brackets at the beginning of the description. If a match is found, the operator type is extracted and returned. Otherwise, a ValueError is raised with a specific error message.

class NodeVizStyleGuide(style='default', print_limit=100)

A class to provide a standardized way to visualize nodes in a graph, particularly for use with graph visualization tools like Graphviz.

style: A string that defines the style of the visualization. Default is ‘default’. print_limit: An integer that sets the maximum number of characters to print for node descriptions and content. Default is 100.

Initialization

Initialize the NodeVizStyleGuide with a specified style and print limit.

style: A string defining the style of the visualization. Default is ‘default’. print_limit: An integer setting the maximum number of characters to print for node descriptions and content. Default is 100.

get_attrs(x)

Get the attributes for a node based on the style guide.

x: The node for which attributes are to be generated.

The get_attrs method takes a node x as input and returns a dictionary of attributes for the node. The attributes include the label, shape, fill color, and style of the node, which are determined based on the node’s properties and the style guide. The method calls other helper methods to construct the label, determine the node shape, assign a color, and set the style.

get_label(x)

Construct a label for a node based on its name, description, and content.

x: The node for which the label is to be constructed.

Using a colon in the name can cause problems in graph visualization tools like Graphviz. To avoid issues, the label is constructed by combining the node’s Python name, truncated description, and content. If the description or content exceeds the print limit, it is truncated and appended with an ellipsis.

get_node_shape(x)

Determine the shape of a node based on its type.

x: The node for which the shape is to be determined.

The shape of a node is determined based on its type. ParameterNode types are represented as ‘box’, while other types are represented as ‘ellipse’.

get_color(x)

Assign a color to a node based on its type.

x: The node for which the color is to be assigned.

The color of a node is determined based on its type. ExceptionNode types are colored ‘firebrick1’, and ParameterNode types are colored ‘lightgray’.

get_style(x)

Set the style of a node based on its properties.

x: The node for which the style is to be set.

The style of a node is set to ‘filled,solid’ if the node is trainable; otherwise, it returns an empty string.

class NodeVizStyleGuideColorful(style='default', print_limit=100)

Bases: opto.trace.nodes.NodeVizStyleGuide

A class to provide a colorful style guide for visualizing nodes in a graph.

style: A string defining the style of the visualization. Default is ‘default’. print_limit: An integer setting the maximum number of characters to print for node descriptions and content. Default is 100.

Initialization

Initialize the NodeVizStyleGuideColorful with a specified style and print limit.

style: A string defining the style of the visualization. Default is ‘default’. print_limit: An integer setting the maximum number of characters to print for node descriptions and content. Default is 100.

get_attrs(x)

Get the attributes for a node based on the colorful style guide.

x: The node for which attributes are to be generated.

The get_attrs method takes a node x as input and returns a dictionary of attributes for the node. The attributes include the label, shape, fill color, style, border color, and border width of the node, which are determined based on the node’s properties and the style guide. The method calls other helper methods to construct the label, determine the node shape, assign a color, and set the style.

get_border_color(x)

Assign a border color to a node based on its type.

x: The node for which the border color is to be assigned.

The border color of a node is determined based on its type. ExceptionNode types are colored ‘firebrick1’, and ParameterNode types are colored ‘black’.

get_color(x)

Assign a fill color to a node based on its type.

x: The node for which the fill color is to be assigned.

The fill color of a node is determined based on its type. ExceptionNode types are colored ‘firebrick1’, and ParameterNode types are colored ‘lightgray’.

get_style(x)

Set the style of a node always as if it is trainable.

class Node(value: Any, *, name: str = None, trainable: bool = False, description: str = '[Node] This is a node in a computational graph.', constraint: Union[None, str] = None, info: Union[None, Dict] = None)

Bases: opto.trace.nodes.AbstractNode[opto.trace.nodes.T]

A data node in a directed graph, this is a basic data structure of Trace.

trainable: A boolean indicating whether the node is trainable or not. _feedback: A dictionary of feedback from children nodes. _description: A string describing the node. _constraint: A string describing all constraints that the data in the node should satisfy. _backwarded: A boolean indicating whether the backward method has been called. _info: A dictionary containing additional information about the node. _dependencies: A dictionary of dependencies on parameters and expandable nodes.

The Node class extends the AbstractNode class to represent a data node in a directed graph. It includes additional attributes and methods to handle feedback, constraints, and dependencies. The node can be marked as trainable, and it can store feedback from children nodes. The node has a description and additional information associated with it. The node can also track dependencies on parameters and expandable nodes, which are nodes that depend on parameters not visible in the current graph level.

The Node class is meant to be subclassed and extended to create specific types of nodes. The feedback mechanism is analogous to gradients in machine learning and is used to propagate information back through the graph. The feedback mechanism is designed to support non-commutative aggregation, so feedback should be handled carefully to maintain the correct order of operations.

Initialization

Initialize an instance of the Node class.

value: The value to be assigned to the node. name: The name of the node (optional). trainable: A boolean indicating whether the node is trainable or not (optional). description: A string describing the node (optional). constraint: A string describing constraints on the node (optional). info: A dictionary containing additional information about the node (optional).

zero_feedback()

Zero out the feedback of the node. zero_feedback should be used judiciously within the feedback propagation process to avoid unintended loss of feedback data. It is specifically designed to be used after feedback has been successfully propagated to parent nodes.

property feedback

The feedback from children nodes.

property description

A textual description of the node.

property info

Additional information about the node.

property type

The type of the data stored in the node.

property parameter_dependencies

The depended parameters.

Ensure that the ‘_dependencies’ attribute is properly initialized and contains a ‘parameter’ key with a corresponding value before calling the parameter_dependencies function to avoid potential KeyError exceptions.

property expandable_dependencies

The depended expandable nodes, where expandable nodes are those who depend on parameters not visible in the current graph level.

Ensure that the ‘_dependencies’ attribute is properly initialized and contains an ‘expandable’ key with a corresponding value before calling the expandable_dependencies function to avoid potential KeyError exceptions

backward(feedback: Any = '', propagator=None, retain_graph=False, visualize=False, simple_visualization=True, reverse_plot=False, print_limit=100)

Performs a backward pass in a computational graph. This function propagates feedback from the current node to its parents, updates the graph visualization if required, and returns the resulting graph.

feedback: The feedback given to the current node. propagator: A function that takes in a node and a feedback, and returns a dict of {parent: parent_feedback}. If not provided, a default GraphPropagator object is used. retain_graph: If True, the graph will be retained after backward pass. visualize: If True, the graph will be visualized using graphviz. simple_visualization: If True, identity operators will be skipped in the visualization; otherwise, they will be included. reverse_plot: if True, plot the graph in reverse order (from child to parent). print_limit: The maximum number of characters to print for node descriptions and content.

The function checks if the current node has already been backwarded. If it has, an AttributeError is raised. Otherwise, the function adds the feedback to the current node by calling the _add_feedback method of the node object. The feedback is initialized with a special “FEEDBACK_ORACLE” node and the propagated feedback from the propagator object.

If the current node has no parents, indicating that it is a root node, the function checks if visualization is enabled. If it is, the current node is added to the digraph object with the appropriate style attributes. Finally, the function returns the digraph object.

If the current node has parents, indicating that it is not a root node, the function initializes a priority queue. The priority queue is used to process the nodes in the correct order during the backward pass. The function enters a loop that continues until the queue is empty. In each iteration, a node is popped from the queue and processed. The node is checked to ensure it has parents and is an instance of the MessageNode class. If not, an AttributeError is raised.

The function propagates information from the current node to its parents by calling the propagator object with the current node as the argument. The propagator object computes the propagated feedback based on the child node’s description, data, and feedback. The propagated feedback is then added to the parents of the current node by calling the _add_feedback method of each parent node.

After processing the parents of the current node, the _backwarded attribute of the current node is updated to indicate that it has been backwarded. This attribute is set to True unless the retain_graph parameter is set to True.

The loop continues until the queue is empty, indicating that all the nodes have been processed. Finally, the function returns the digraph object.

clone()

Create and return a duplicate of the current Node object.

detach()

Create and return a deep copy of the current instance of the Node class.

getattr(key)

Get the attribute of the node with the specified key.

key: The key of the attribute to get.

call(fun: str, *args, **kwargs)

Call the function with the specified arguments and keyword arguments.

fun: The function to call. args: The arguments to pass to the function. kwargs: The keyword arguments to pass to the function.

len()

Return the length of the node.

We overload magic methods that return a value. This method returns a MessageNode.

eq(other)

Check if the node is equal to another value.

other: The value to compare the node to.

If a logic operator is used in an if-statement, it will return a boolean value. Otherwise, it will return a MessageNode.

neq(other)

Check if the node is not equal to another value.

other: The value to compare the node to.

If a logic operator is used in an if-statement, it will return a boolean value. Otherwise, it will return a MessageNode.

format(*args, **kwargs)

capitalize()

lower()

upper()

swapcase()

title()

split(sep=None, maxsplit=-1)

strip(chars=None)

replace(old, new, count=-1)

items()

values()

keys()

pop(__index=-1)

append(*args, **kwargs)

class ParameterNode(value, *, name=None, trainable=True, description='[ParameterNode] This is a ParameterNode in a computational graph.', constraint=None, info=None)

Bases: opto.trace.nodes.Node[opto.trace.nodes.T]

class MessageNode(value, *, inputs: Union[List[opto.trace.nodes.Node], Dict[str, opto.trace.nodes.Node]], description: str, constraint=None, name=None, info=None)

Bases: opto.trace.nodes.Node[opto.trace.nodes.T]

Output of an operator.

description: a string to describe the operator it begins with [operator_name] and then describes the operator. When referring to inputs use the keys in args (if args is a dict), or the names of the nodes in args (if args is a list). Here’re some examples:

MessageNode(node_a, inputs=[node_a], description=”[identity] This is an identity operator.”) MessageNode(copy_node_a, inputs=[node_a], description=”[copy] This is a copy operator.”) MesssageNode(1, inputs={‘a’:node_a, ‘b’:node_b}, description=”[Add] This is an add operator of a and b.”)

Initialization

Initialize an instance of the Node class.

value: The value to be assigned to the node. name: The name of the node (optional). trainable: A boolean indicating whether the node is trainable or not (optional). description: A string describing the node (optional). constraint: A string describing constraints on the node (optional). info: A dictionary containing additional information about the node (optional).

property inputs

property hidden_dependencies

Returns the set of hidden dependencies that are not visible in the current graph level.

class ExceptionNode(value: Exception, *, inputs: Union[List[opto.trace.nodes.Node], Dict[str, opto.trace.nodes.Node]], description: str = '[ExceptionNode] This is node containing the error of execution.', constraint=None, name=None, info=None)

Bases: opto.trace.nodes.MessageNode[opto.trace.nodes.T]

Node containing the exception message.

Initialization

Initialize an instance of the Node class.

value: The value to be assigned to the node. name: The name of the node (optional). trainable: A boolean indicating whether the node is trainable or not (optional). description: A string describing the node (optional). constraint: A string describing constraints on the node (optional). info: A dictionary containing additional information about the node (optional).

create_feedback(style='simple')