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