tinygrad graph

Note

You likely want the upstream tinygrad, not tinygrab. Tinygrab contains AI generated docstrings for a tinygrad snapshot. Upstream: https://tinygrad.org

tinygrad.graph.add_st_node(nmx, nmo, label, st: ShapeTracker)

tinygrad.graph.get_sop(op: List[UnaryOps | BinaryOps | ReduceOps | MovementOps | LoadOps | TernaryOps | BufferOps])

Returns a string representation of a list of operations ‘op’. The returned string is based on the last part of the class name of each operation in the list, after splitting by ‘.’. If the length of ‘op’ is less than or equal to 2, it returns the full names of all operations in reverse order, separated by ‘.’. If the length is between 3 and 6, it returns only the first three characters of each operation name. If the length of ‘op’ is more than 6, it simply returns the string representation of the length of ‘op’.

Parameters:

op (List[Op]) – List of operations for which a string representation will be returned.

Returns:

String representation of the list of operations ‘op’

tinygrad.graph.graph_uops(uops: List[UOp])

This function generates a directed graph from a list of UOps. It creates a node for each UOp and adds edges between them based on their input/output relationships. The nodes are colored according to the type of UOp they represent.

Parameters:

uops – A list of UOp objects representing the operations to be graphed.

tinygrad.graph.log_schedule_item(si: ScheduleItem)

Log the schedule item in the graph.

This function logs the schedule item, its operations, and input/output relationships in a directed graph. The graph is used for visualization purposes and to keep track of the transformations on the data.

Parameters:

si – ScheduleItem object representing an operation or transformation.

tinygrad.graph.nm(x)

Assign a unique node_id to an object x if it doesn’t have one already.

This function checks for the presence of ‘node_id’ attribute in the given object ‘x’. If not found, it assigns a new unique id by incrementing the global ‘node_count’ variable and associates this id with the object ‘x’. Finally, it returns the node_id of the object ‘x’.

Parameters:

x (Any) – The input object to which a unique node_id will be assigned if not already present.

Returns:

The node_id of the object ‘x’

tinygrad.graph.print_tree(lazydata: LazyOp)

This function takes a single argument of type LazyOp and prints the tree structure.

Parameters:

lazydata (LazyOp) – The input data for which we want to generate the tree structure.

Returns:

None. The function prints the tree structure directly.

tinygrad.graph.str_dtype(dtyp)

Returns a string representation of the type ‘dtyp’. If ‘dtyp’ is of type ‘float’, it returns an empty string. Otherwise, it returns the string representation of ‘dtyp’.

Parameters:

dtyp (Any) – The input object whose string representation will be returned.

Returns:

String representation of the object ‘dtyp’ or an empty string if ‘dtyp’ is of type ‘float’.