tinygrad shape.symbolic

Note

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

class tinygrad.shape.symbolic.AndNode(nodes: List[Node])[source]

Bases: RedNode

AndNode class for creating AND operations on a list of nodes.

RedNode

Parent class for this class.

Type:: Class

substitute(var_vals: Dict[Variable | NumNode, Node]) → Node[source]

Substitutes variables in the AND operation with their corresponding values provided in var_vals.

Parameters:

self (AndNode) – The instance of AndNode.
var_vals (Dict[VariableOrNum, Node]) – A dictionary containing variable-value pairs for substitution.

Returns:

A new node with variables substituted by their corresponding values.

Return type:

Node

class tinygrad.shape.symbolic.DivNode(a: Node, b: Node | int)[source]

Bases: OpNode

The DivNode class is a subclass of OpNode. It defines the behavior of division operations on nodes.

a

The dividend node.

Type:: Node

b

The divisor integer.

Type:: int

get_bounds() → Tuple[int, int][source]

Get the lower and upper bounds of a DivNode.

Returns:: A tuple with the lower and upper bounds.

substitute(var_vals: Dict[Variable | NumNode, Node]) → Node[source]

Substitute the variables in a DivNode with their corresponding nodes.

Parameters:: var_vals – A dictionary containing the variables as keys and the corresponding nodes as values.
Returns:: A new DivNode with the variables substituted by their nodes.

class tinygrad.shape.symbolic.LtNode(a: Node, b: Node | int)[source]

Bases: OpNode

This class represents the lower than operator for nodes.

a

The first node to compare.

Type:: Node

b

The second node or number to compare.

Type:: Union[Node, int]

get_bounds() → Tuple[int, int][source]

Get the bounds of the LtNode.

Parameters:: self (LtNode) – The instance of LtNode.
Returns:: A tuple representing the lower and upper bounds.
Return type:: Tuple[int, int]

substitute(var_vals: Dict[Variable | NumNode, Node]) → Node[source]

Substitute variables in the LtNode with their corresponding values.

Parameters:

self (LtNode) – The instance of LtNode.
var_vals (Dict[VariableOrNum, Node]) – A dictionary mapping variables to their values.

Returns:

A new node with the variables substituted by their corresponding values.

Return type:

Node

class tinygrad.shape.symbolic.ModNode(a: Node, b: Node | int)[source]

Bases: OpNode

The ModNode class is a subclass of OpNode. It defines the behavior of modulo operations on nodes.

a

The dividend node.

Type:: Node

b

The divisor integer.

Type:: int

get_bounds() → Tuple[int, int][source]

Get the lower and upper bounds of a ModNode.

Returns:: A tuple with the lower and upper bounds.

substitute(var_vals: Dict[Variable | NumNode, Node]) → Node[source]

Substitute the variables in a ModNode with their corresponding nodes.

Parameters:: var_vals – A dictionary containing the variables as keys and the corresponding nodes as values.
Returns:: A new ModNode with the variables substituted by their nodes.

class tinygrad.shape.symbolic.MulNode(a: Node, b: Node | int)[source]

Bases: OpNode

This class represents a multiplication node.

a

The first operand in the multiplication operation.

Type:: Node

b

The second operand in the multiplication operation.

Type:: Union[Node, int]

get_bounds() → Tuple[int, int][source]

This method returns the lower and upper bounds of the current node.

Returns:: A tuple containing the lower and upper bounds.
Return type:: Tuple[int, int]

substitute(var_vals: Dict[Variable | NumNode, Node]) → Node[source]

This method substitutes variables in the current node with their corresponding values.

Parameters:: var_vals (Dict[VariableOrNum, Node]) – A dictionary containing variable-value pairs.
Returns:: A new node with variables replaced by their corresponding values.
Return type:: Node

class tinygrad.shape.symbolic.Node[source]

Bases: object

The base class for a node in the expression tree. This class represents a single point in the grid of possible expressions.

b

The base node or integer value.

Type:: Union[Node, int]

min

The minimum value of this node.

Type:: int

max

The maximum value of this node.

Type:: int

static ands(nodes: List[Node]) → Node[source]

This function takes a list of nodes as input and returns a single node.

nodes

A list of nodes to be processed.

Type:: List[Node]

Returns:: A single node that represents the logical AND of all input nodes. If any node in the input list is False, it will return a NumNode(0). If all nodes are True, it will return the last node in the list. If there’s only one unique node (excluding NumNode(1)), it will return this node.
Return type:: Node

b: Node | int

expand(idxs: Tuple[Variable | NumNode, ...] | None = None) → List[Node][source]

Expand a Node into List[Node] that enumerates the underlying Variables from min to max.

The expansion increments earlier variables faster than later variables (as specified in the argument).

Parameters:: idxs (Optional[Tuple[VariableOrNum, ...]]) – The indices to expand. Defaults to None.
Returns:: The expanded nodes.
Return type:: List[Node]

expand_idx() → Variable | NumNode[source]

Expand a Node into a single Variable or an integer if the underlying Variables are not defined.

Returns:: The expanded node.
Return type:: VariableOrNum

property hash: int

Returns a hash value for the node based on its key.

Returns:: A hash value of the node’s key.
Return type:: int

static iter_idxs(idxs: Tuple[Variable | NumNode, ...]) → Iterator[Tuple[int, ...]][source]

Get an iterator over the indices from min to max for each Variable.

Parameters:: idxs (Tuple[VariableOrNum, ...]) – The indices to iterate over.
Returns:: An iterator over the indices.
Return type:: Iterator[Tuple[int, …]]

property key: str

Returns a string representation of the node.

ctx

The context for rendering the node. Defaults to “DEBUG”.

Type:: str

Returns:: A string representation of the node.
Return type:: str

max: int

min: int

render(ops=None, ctx=None) → Any[source]

Render the expression tree as a string or other object depending on the output type.

Parameters:

ops (Optional[Dict[Type[Node], Callable]]) – The dictionary of rendering functions for each node type. Defaults to render_python.
ctx (Optional[Any]) – The context for the rendering function.

Returns:

The rendered object.

Return type:

Any

substitute(var_vals: Dict[Variable | NumNode, Node]) → Node[source]

Substitute Variables with the values in var_vals.

Parameters:: var_vals (Dict[VariableOrNum, Node]) – The dictionary of variable substitutions.
Returns:: The node with variables substituted.
Return type:: Node

static sum(nodes: List[Node]) → Node[source]

Calculate the sum of nodes with max or min attributes.

Parameters:: nodes (List[Node]) – Input list of nodes to calculate sum for.
Returns:: Summed node calculated from input nodes.
Return type:: Node

unbind() → Tuple[Node, int | None][source]

Unbind the node by replacing any bound Variables with their unbound counterparts.

Returns:: The unbound node and None.
Return type:: Tuple[Node, Optional[int]]

vars() → Set[Variable][source]

Get the set of variables in this node.

Returns:: The set of variables.
Return type:: Set[Variable]

class tinygrad.shape.symbolic.NumNode(num: int)[source]

Bases: Node

Node class to represent a number.

b

The numeric value of the node.

Type:: int

min

Minimum possible value for this node.

Type:: int

max

Maximum possible value for this node.

Type:: int

bind(val)[source]

Bind the value of this node to another value.

Parameters:: val – The value to bind this node to.
Returns:: This node with its value bound to val.
Return type:: NumNode
Raises:: AssertionError – If self.b is not equal to val.

substitute(var_vals: Dict[Variable | NumNode, Node]) → Node[source]

Substitute variables in this node with their corresponding values.

Parameters:: var_vals (Dict[VariableOrNum, Node]) – A dictionary mapping variables to their substituted values.
Returns:: This node with all its variables substituted by their values.
Return type:: NumNode

class tinygrad.shape.symbolic.OpNode(a: Node, b: Node | int)[source]

Bases: Node

An operation node class that represents a node with an operation involving two nodes or a node and an integer value.

a

The first input node.

Type:: Node

b

The second input node or an integer value.

Type:: Union[Node, int]

get_bounds() → Tuple[int, int][source]

Get the bounds of this operation node. This method should be implemented in any child classes.

Raises:: NotImplementedError – If this method is not overridden by a child class.

vars()[source]

Get the set of variables involved in this operation node.

Returns:: A set of nodes representing the variables involved in this operation.
Return type:: Set[Node]

class tinygrad.shape.symbolic.RedNode(nodes: List[Node])[source]

Bases: Node

Base class for all “red” nodes.

nodes

The list of child nodes.

Type:: List[Node]

vars() → Set[Variable][source]

Returns the union of all variables used in the child nodes.

Returns:: A set containing all unique variables used by the child nodes.
Return type:: Set[Variable]

class tinygrad.shape.symbolic.SumNode(nodes: List[Node])[source]

Bases: RedNode

SumNode class for representing sum of nodes.

__mul__[source]

Distribute multiplication over summation.

Type:: method

__floordiv__[source]

Perform floor division with consideration to factoring.

Type:: method

property flat_components

Get the flattened components of a SumNode.

If a component is an instance of SumNode, it will be expanded; otherwise, it remains unchanged. This method is used to flatten a nested sum of nodes into a flat list of nodes.

Returns:: The flattened components of the SumNode.
Return type:: List[Node]

substitute(var_vals: Dict[Variable | NumNode, Node]) → Node[source]

Substitute the variables in the expression with their corresponding values.

Parameters:: var_vals (Dict[VariableOrNum, Node]) – A dictionary containing variable-value pairs.
Returns:: The resulting node after substitution.
Return type:: Node

class tinygrad.shape.symbolic.Variable(expr: str | None, nmin: int, nmax: int)[source]

Bases: Node

Variable class.

expr

Optional string representation of the variable.

Type:: Optional[str]

nmin

Minimum value for the variable.

Type:: int

nmax

Maximum value for the variable.

Type:: int

bind(val)[source]

Bind the Variable to a specific value.

Parameters:

self – Self reference.
val (int) – The value to bind the Variable to.

Returns:

The bound Variable.

Return type:

Variable

Raises:

AssertionError – If the Variable is already bound or the value is outside the specified range.

substitute(var_vals: Dict[Variable | NumNode, Node]) → Node[source]

Substitute the Variable with its value or return itself if not in var_vals.

Parameters:

self – Self reference.
var_vals (Dict[VariableOrNum, Node]) – A dictionary mapping Variables to their values.

Returns:

The substituted Node.

Return type:

Node

unbind() → Tuple[Variable, int][source]

Unbind the Variable from its value and return a new Variable with the same properties and its value.

Parameters:: self – Self reference.
Returns:: A tuple containing a new Variable object and the value it was bound to.
Return type:: Tuple[Variable, int]
Raises:: AssertionError – If the Variable is not bound.

property val

Get the value of the Variable.

Parameters:: self – Self reference.
Returns:: The value of the Variable.
Return type:: int
Raises:: AssertionError – If the Variable is not bound to a value.

vars()[source]

Get a set of the Variable.

Parameters:: self – Self reference.
Returns:: A set containing the Variable.
Return type:: Set[Variable]

tinygrad.shape.symbolic.create_node(ret: Node)[source]

Create a node object based on the provided parameters.

Parameters:: ret (Node) – The node to be created.
Returns:: If min is equal to max, return NumNode with value min. Else, return the inputted node.
Return type:: Union[NumNode, Node]
Raises:: AssertionError – If minimum value is greater than maximum value.

tinygrad.shape.symbolic.create_rednode(typ: Type[RedNode], nodes: List[Node])[source]

Create a new red node.

Parameters:

typ (Type[RedNode]) – The type of the red node to be created.
nodes (List[Node]) – The list of nodes to be used for creation.

Returns:

A new red node created with given type and nodes.

tinygrad.shape.symbolic.is_sym_int(x: Any) → bool[source]

Check if an object is either of type int or an instance of the Node class.

tinygrad.shape.symbolic.x

The object to check.

Type:: Any

Returns:: True if x is of type int or an instance of Node, False otherwise.
Return type:: bool

tinygrad.shape.symbolic.sym_infer(a: Node | int, var_vals: Dict[Variable, int]) → int[source]

Perform inference on the given node or integer with the provided variable values.

Parameters:

a (Union[Node, int]) – The node or integer to perform inference on.
var_vals (Dict[Variable, int]) – A dictionary mapping variables to their respective integer values.

Returns:

An integer result of the inference.

Return type:

int

tinygrad.shape.symbolic.sym_render(a: Node | int, ops=None, ctx=None) → str[source]

Render a symbolic representation of the given node or integer.

Parameters:

a (Union[Node, int]) – The node or integer to be rendered.
ops (Any) – Optional operations for rendering. Default is None.
ctx (Any) – Optional context for rendering. Default is None.

Returns:

A string representation of the given node or integer.

Return type:

str