pf.compute_graph

Submodules

• pf.compute_graph.compute_graph — procfunc.compute_graph.compute_graph

• pf.compute_graph.node — procfunc.compute_graph.node

• pf.compute_graph.operators_info — procfunc.compute_graph.operators_info

• pf.compute_graph.proxy — procfunc.compute_graph.proxy

• pf.compute_graph.util — procfunc.compute_graph.util

Members

class procfunc.compute_graph.ComputeGraph

Bases: object

ComputeGraph(inputs: procfunc.util.pytree.PyTree[typing.Any, procfunc.compute_graph.node.Node], outputs: procfunc.util.pytree.PyTree[typing.Any, procfunc.compute_graph.node.Node], name: str, metadata: dict[str, typing.Any])

inputs: PyTree[Any, Node]

outputs: PyTree[Any, Node]

name: str

metadata: dict[str, Any]

clear_values()

__init__(inputs, outputs, name, metadata)

Parameters:

• inputs (PyTree[Any, Node])

• outputs (PyTree[Any, Node])

• name (str)

• metadata (dict[str, Any])

Return type:

None

Supported special methods: __call__

class procfunc.compute_graph.ConstantNode

Bases: Node

__init__(value, metadata=None)

Parameters:

• value (Any)

• metadata (dict[str, Any])

class procfunc.compute_graph.FunctionCallNode

Bases: Node

__init__(func, args, kwargs, metadata=None)

Parameters:

• func (Callable[[...], Any])

• args (tuple)

• kwargs (dict)

• metadata (dict[str, Any])

class procfunc.compute_graph.GetAttributeNode

Bases: Node

GetAttributeNode(source: procfunc.compute_graph.node.Node, attribute_name: str, metadata: dict[str, typing.Any] = None)

__init__(source, attribute_name, metadata=None)

Parameters:

• source (Node)

• attribute_name (str)

• metadata (dict[str, Any])

class procfunc.compute_graph.InputPlaceholderNode

Bases: Node

__init__(name, default_value, metadata=None)

Parameters:

• name (str)

• default_value (Any)

• metadata (dict[str, Any])

class procfunc.compute_graph.MethodCallNode

Bases: Node

represents an {args[0]}.{method_name}(*args[1:], **kwargs) call

• the node to be used as self is the first arg, since it is a dynamic value

• the method name is assumed to be const

__init__(callee, method_name, args, kwargs, metadata=None)

Parameters:

• callee (Node)

• method_name (str)

• args (tuple)

• kwargs (dict)

• metadata (dict[str, Any])

class procfunc.compute_graph.MutatedArgumentNode

Bases: Node

__init__(original_node, mutator_call_node, metadata=None)

Parameters:

• original_node (Node)

• mutator_call_node (FunctionCallNode | MethodCallNode)

• metadata (dict[str, Any])

class procfunc.compute_graph.Node

Bases: object

__init__(args, kwargs, metadata=None)

Parameters:

• args (tuple)

• kwargs (dict)

• metadata (dict[str, Any])

inputs_pytree()

Return type:

PyTree

class procfunc.compute_graph.ProceduralNode

Bases: Node

ProceduralNode(node_type: str, attrs: dict[str, typing.Any], kwargs: dict, metadata: dict[str, typing.Any] = None)

__init__(node_type, attrs, kwargs, metadata=None)

Parameters:

• node_type (str)

• attrs (dict[str, Any])

• kwargs (dict)

• metadata (dict[str, Any])

class procfunc.compute_graph.SubgraphCallNode

Bases: Node

__init__(subgraph, args, kwargs, metadata=None)

Parameters:

• subgraph (ComputeGraph)

• args (tuple)

• kwargs (dict)

• metadata (dict[str, Any])

procfunc.compute_graph.normalize_args_to_kwargs(func, args, kwargs)

Try to fully populate kwargs, by moving over positional args & filling in defaults

Some args may not be able to be converted to kwargs, e.g. *args have no names that work

Parameters:

• func (Callable) – The function whose signature we should respect

• args (tuple) – The original positional arguments to the function

• kwargs (dict) – The keyword arguments to the function

Returns:

A tuple of (args, kwargs) where args is a tuple of positional arguments and kwargs is a dictionary of keyword arguments.

Return type:

tuple[tuple, dict]

GUARANTEE: func(*returned_args, **returned_kwargs) == func(*args, **kwargs) and does not crash

class procfunc.compute_graph.OperatorType

Bases: Enum

ADD = 'add'

SUB = 'sub'

MUL = 'mul'

DIV = 'div'

TRUEDIV = 'truediv'

POW = 'pow'

MOD = 'mod'

LESS_THAN = 'lt'

LESS_THAN_EQUAL = 'le'

GREATER_THAN = 'gt'

GREATER_THAN_EQUAL = 'ge'

EQUAL = 'eq'

NOT_EQUAL = 'ne'

AND = 'and'

OR = 'or'

NOT = 'invert'

LSHIFT = 'lshift'

RSHIFT = 'rshift'

BIT_AND = 'and'

BIT_OR = 'or'

BIT_XOR = 'xor'

GETITEM = 'getitem'

VECTOR_PACK = 'VECTOR_PACK'

NOOP = 'NOOP'

Supported special methods: __contains__, __getitem__, __iter__, __len__

class procfunc.compute_graph.AttributeProxy

Bases: Proxy

Special proxy for attribute access that supports peekthrough optimization

__init__(node)

Parameters:

node (Node)

Supported special methods: __abs__, __add__, __and__, __bool__, __call__, __floordiv__, __getattr__, __getitem__, __invert__, __iter__, __len__, __lshift__, __mod__, __mul__, __neg__, __or__, __pos__, __pow__, __radd__, __rand___, __rfloordiv__, __rlshift__, __rmod__, __rmul__, __ror___, __rpow__, __rrshift__, __rshift__, __rsub__, __rtruediv__, __rxor__, __sub__, __truediv__, __xor__

class procfunc.compute_graph.Proxy

Bases: Generic[T]

General-purpose wrapper for Node that provides all dunder methods.

node: Node

__init__(node)

Parameters:

node (Node)

Return type:

None

Supported special methods: __abs__, __add__, __and__, __bool__, __call__, __floordiv__, __getattr__, __getitem__, __invert__, __iter__, __len__, __lshift__, __mod__, __mul__, __neg__, __or__, __pos__, __pow__, __radd__, __rand___, __rfloordiv__, __rlshift__, __rmod__, __rmul__, __ror___, __rpow__, __rrshift__, __rshift__, __rsub__, __rtruediv__, __rxor__, __sub__, __truediv__, __xor__

class procfunc.compute_graph.LiteralConstant

Bases: object

__init__(value)

Parameters:

value (Any)

procfunc.compute_graph.graph_nodes_equal(graph1, graph2)

Parameters:

• graph1 (ComputeGraph)

• graph2 (ComputeGraph)

Return type:

bool

procfunc.compute_graph.transform_compute_graph(compute_graph, transform_fn, graph_name=None)

Parameters:

• compute_graph (ComputeGraph)

• transform_fn (Callable[[Node], Any])

• graph_name (str | None)

procfunc.compute_graph.transform_nodetree(root, transform_fn, memo={})

Parameters:

• root (Node)

• transform_fn (Callable[[Node], Any])

• memo (dict[int, Node])

procfunc.compute_graph.traverse_breadth_first(graph, yield_parent=False, yield_name=False, yield_consts=False)

Traverse all nodes in the compute graph.

Parameters:

• graph (ComputeGraph) – The compute graph to traverse

• yield_parent (bool) – If True, yield (parent, child), with output nodes having parent=None

• yield_name (bool) – If True, yield (name, child) or (name, parent, child) if yield_parent is also True

• yield_consts (bool) – If True, yield child arguments of nodes even if they are not Nodes

Return type:

Generator[Any, None, None]

procfunc.compute_graph.traverse_depth_first(graph, yield_parent=False, yield_name=False, yield_consts=False, order='postorder')

Parameters:

• graph (ComputeGraph)

• yield_parent (bool)

• yield_name (bool)

• yield_consts (bool)

• order (Literal['preorder', 'postorder'])

Return type:

Generator[Any, None, None]

procfunc.compute_graph.traverse_depth_first_node(node, yield_consts=False, order='postorder')

Parameters:

• node (Node)

• yield_consts (bool)

• order (Literal['preorder', 'postorder'])

Return type:

Generator[Any, None, None]

procfunc.compute_graph.traverse_nested_graphs(graph, yield_call_nodes=False)

Parameters:

• graph (ComputeGraph)

• yield_call_nodes (bool)

Return type:

Generator[tuple[Node | None, ComputeGraph], None, None]

procfunc.compute_graph.usages_per_node(graph)

Parameters:

graph (ComputeGraph)

Return type:

dict[int, list[Node]]

pf.compute_graph.compute_graph

class procfunc.compute_graph.compute_graph.ComputeGraph

Bases: object

ComputeGraph(inputs: procfunc.util.pytree.PyTree[typing.Any, procfunc.compute_graph.node.Node], outputs: procfunc.util.pytree.PyTree[typing.Any, procfunc.compute_graph.node.Node], name: str, metadata: dict[str, typing.Any])

inputs: PyTree[Any, Node]

outputs: PyTree[Any, Node]

name: str

metadata: dict[str, Any]

clear_values()

__init__(inputs, outputs, name, metadata)

Parameters:

• inputs (PyTree[Any, Node])

• outputs (PyTree[Any, Node])

• name (str)

• metadata (dict[str, Any])

Return type:

None

Supported special methods: __call__

pf.compute_graph.node

class procfunc.compute_graph.node.Node

Bases: object

__init__(args, kwargs, metadata=None)

Parameters:

• args (tuple)

• kwargs (dict)

• metadata (dict[str, Any])

inputs_pytree()

Return type:

PyTree

class procfunc.compute_graph.node.SubgraphCallNode

Bases: Node

__init__(subgraph, args, kwargs, metadata=None)

Parameters:

• subgraph (ComputeGraph)

• args (tuple)

• kwargs (dict)

• metadata (dict[str, Any])

class procfunc.compute_graph.node.FunctionCallNode

Bases: Node

__init__(func, args, kwargs, metadata=None)

Parameters:

• func (Callable[[...], Any])

• args (tuple)

• kwargs (dict)

• metadata (dict[str, Any])

class procfunc.compute_graph.node.MethodCallNode

Bases: Node

represents an {args[0]}.{method_name}(*args[1:], **kwargs) call

• the node to be used as self is the first arg, since it is a dynamic value

• the method name is assumed to be const

__init__(callee, method_name, args, kwargs, metadata=None)

Parameters:

• callee (Node)

• method_name (str)

• args (tuple)

• kwargs (dict)

• metadata (dict[str, Any])

class procfunc.compute_graph.node.GetAttributeNode

Bases: Node

GetAttributeNode(source: procfunc.compute_graph.node.Node, attribute_name: str, metadata: dict[str, typing.Any] = None)

__init__(source, attribute_name, metadata=None)

Parameters:

• source (Node)

• attribute_name (str)

• metadata (dict[str, Any])

class procfunc.compute_graph.node.ProceduralNode

Bases: Node

ProceduralNode(node_type: str, attrs: dict[str, typing.Any], kwargs: dict, metadata: dict[str, typing.Any] = None)

__init__(node_type, attrs, kwargs, metadata=None)

Parameters:

• node_type (str)

• attrs (dict[str, Any])

• kwargs (dict)

• metadata (dict[str, Any])

class procfunc.compute_graph.node.MutatedArgumentNode

Bases: Node

__init__(original_node, mutator_call_node, metadata=None)

Parameters:

• original_node (Node)

• mutator_call_node (FunctionCallNode | MethodCallNode)

• metadata (dict[str, Any])

class procfunc.compute_graph.node.ConstantNode

Bases: Node

__init__(value, metadata=None)

Parameters:

• value (Any)

• metadata (dict[str, Any])

class procfunc.compute_graph.node.InputPlaceholderNode

Bases: Node

__init__(name, default_value, metadata=None)

Parameters:

• name (str)

• default_value (Any)

• metadata (dict[str, Any])

procfunc.compute_graph.node.normalize_args_to_kwargs(func, args, kwargs)

Try to fully populate kwargs, by moving over positional args & filling in defaults

Some args may not be able to be converted to kwargs, e.g. *args have no names that work

Parameters:

• func (Callable) – The function whose signature we should respect

• args (tuple) – The original positional arguments to the function

• kwargs (dict) – The keyword arguments to the function

Returns:

A tuple of (args, kwargs) where args is a tuple of positional arguments and kwargs is a dictionary of keyword arguments.

Return type:

tuple[tuple, dict]

GUARANTEE: func(*returned_args, **returned_kwargs) == func(*args, **kwargs) and does not crash

pf.compute_graph.operators_info

class procfunc.compute_graph.operators_info.OperatorType

Bases: Enum

ADD = 'add'

SUB = 'sub'

MUL = 'mul'

DIV = 'div'

TRUEDIV = 'truediv'

POW = 'pow'

MOD = 'mod'

LESS_THAN = 'lt'

LESS_THAN_EQUAL = 'le'

GREATER_THAN = 'gt'

GREATER_THAN_EQUAL = 'ge'

EQUAL = 'eq'

NOT_EQUAL = 'ne'

AND = 'and'

OR = 'or'

NOT = 'invert'

LSHIFT = 'lshift'

RSHIFT = 'rshift'

BIT_AND = 'and'

BIT_OR = 'or'

BIT_XOR = 'xor'

GETITEM = 'getitem'

VECTOR_PACK = 'VECTOR_PACK'

NOOP = 'NOOP'

Supported special methods: __contains__, __getitem__, __iter__, __len__

pf.compute_graph.proxy

General-purpose Proxy wrapper for Node with all dunders.

class procfunc.compute_graph.proxy.Proxy

Bases: Generic[T]

General-purpose wrapper for Node that provides all dunder methods.

node: Node

__init__(node)

Parameters:

node (Node)

Return type:

None

Supported special methods: __abs__, __add__, __and__, __bool__, __call__, __floordiv__, __getattr__, __getitem__, __invert__, __iter__, __len__, __lshift__, __mod__, __mul__, __neg__, __or__, __pos__, __pow__, __radd__, __rand___, __rfloordiv__, __rlshift__, __rmod__, __rmul__, __ror___, __rpow__, __rrshift__, __rshift__, __rsub__, __rtruediv__, __rxor__, __sub__, __truediv__, __xor__

class procfunc.compute_graph.proxy.AttributeProxy

Bases: Proxy

Special proxy for attribute access that supports peekthrough optimization

__init__(node)

Parameters:

node (Node)

Supported special methods: __abs__, __add__, __and__, __bool__, __call__, __floordiv__, __getattr__, __getitem__, __invert__, __iter__, __len__, __lshift__, __mod__, __mul__, __neg__, __or__, __pos__, __pow__, __radd__, __rand___, __rfloordiv__, __rlshift__, __rmod__, __rmul__, __ror___, __rpow__, __rrshift__, __rshift__, __rsub__, __rtruediv__, __rxor__, __sub__, __truediv__, __xor__

pf.compute_graph.util

class procfunc.compute_graph.util.LiteralConstant

Bases: object

__init__(value)

Parameters:

value (Any)

procfunc.compute_graph.util.traverse_breadth_first(graph, yield_parent=False, yield_name=False, yield_consts=False)

Traverse all nodes in the compute graph.

Parameters:

• graph (ComputeGraph) – The compute graph to traverse

• yield_parent (bool) – If True, yield (parent, child), with output nodes having parent=None

• yield_name (bool) – If True, yield (name, child) or (name, parent, child) if yield_parent is also True

• yield_consts (bool) – If True, yield child arguments of nodes even if they are not Nodes

Return type:

Generator[Any, None, None]

procfunc.compute_graph.util.traverse_depth_first_node(node, yield_consts=False, order='postorder')

Parameters:

• node (Node)

• yield_consts (bool)

• order (Literal['preorder', 'postorder'])

Return type:

Generator[Any, None, None]

procfunc.compute_graph.util.traverse_depth_first(graph, yield_parent=False, yield_name=False, yield_consts=False, order='postorder')

Parameters:

• graph (ComputeGraph)

• yield_parent (bool)

• yield_name (bool)

• yield_consts (bool)

• order (Literal['preorder', 'postorder'])

Return type:

Generator[Any, None, None]

procfunc.compute_graph.util.traverse_nested_graphs(graph, yield_call_nodes=False)

Parameters:

• graph (ComputeGraph)

• yield_call_nodes (bool)

Return type:

Generator[tuple[Node | None, ComputeGraph], None, None]

procfunc.compute_graph.util.usages_per_node(graph)

Parameters:

graph (ComputeGraph)

Return type:

dict[int, list[Node]]

procfunc.compute_graph.util.graph_nodes_equal(graph1, graph2)

Parameters:

• graph1 (ComputeGraph)

• graph2 (ComputeGraph)

Return type:

bool

procfunc.compute_graph.util.transform_nodetree(root, transform_fn, memo={})

Parameters:

• root (Node)

• transform_fn (Callable[[Node], Any])

• memo (dict[int, Node])

procfunc.compute_graph.util.transform_compute_graph(compute_graph, transform_fn, graph_name=None)

Parameters:

• compute_graph (ComputeGraph)

• transform_fn (Callable[[Node], Any])

• graph_name (str | None)