tinygrad lazy

Note

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

Bases: object

LazyBuffer class for lazy operations on buffers.

__deletable__

Tuple containing the attribute ‘op’ which can be deleted.

Type:: tuple

device

Device where this buffer is located.

Type:: str

st

Shape tracker object.

Type:: ShapeTracker

optype

Operation type.

Type:: OpType

op

Lazy operation object.

Type:: Optional[LazyOp]

dtype

Data type of the buffer.

Type:: DType

src

Source buffer, if any. Default is None.

Type:: Optional[Buffer]

base

Base lazy buffer, if any. Default is None.

Type:: Optional[LazyBuffer]

output_buffer

Output buffer. Default is None.

Type:: Optional[Buffer]

children

Weak set of child lazy buffers.

Type:: WeakSet[LazyBuffer]

views

Weak set of view lazy buffers.

Type:: WeakSet[LazyBuffer]

property base

Return the base of this LazyBuffer.

Returns:: The base of this LazyBuffer if it is not None, else self.

property buffers: Tuple[LazyBuffer, ...]

Return a tuple containing the instance of LazyBuffer.

Returns:: A tuple containing the instance of LazyBuffer.
Return type:: Tuple[LazyBuffer]

cast(dtype: DType, bitcast: bool = False)[source]

Cast the elements of this buffer to a new data type.

dtype

The desired data type for the elements of this buffer.

Type:: DType

bitcast

Whether to allow a bit-level cast, which is faster but may result in unspecified behavior if used incorrectly. Defaults to False.

Type:: bool

Returns:: This method modifies the LazyBuffer object in-place and returns it for chaining purposes.
Return type:: self

const(val: float | int) → LazyBuffer[source]

Creates a new constant LazyBuffer object based on the shape and data type of this instance.

Returns:: A new constant LazyBuffer object with the same shape and data type as this instance.
Return type:: LazyBuffer

contiguous() → LazyBuffer[source]

This method is used to ensure a lazy buffer is stored in a contiguous manner. If the source buffer is already contiguous, it will return itself directly. Otherwise, it will create a new contiguous lazy buffer.

self

The source lazy buffer.

Type:: LazyBuffer

Returns:: A contiguous version of the source lazy buffer.
Return type:: LazyBuffer

copy_to_device(device: str) → LazyBuffer[source]

This method is used to copy a lazy buffer to a specified device. It will check if the buffer is already on the target device and return it directly if true. Otherwise, it will create a new lazy buffer on the target device.

self

The source lazy buffer.

Type:: LazyBuffer

device

The target device to copy the buffer to.

Type:: str

Returns:: The copied lazy buffer on the target device.
Return type:: LazyBuffer

property dtype

Get the data type of this LazyBuffer.

Returns:: The data type of this LazyBuffer.

e(op: UnaryOps | BinaryOps | TernaryOps, *srcs: LazyBuffer, arg: Any | None = None) → LazyBuffer[source]

This method performs an operation on the input buffers and returns a new LazyBuffer.

Parameters:

self (LazyBuffer) – The instance of the LazyBuffer class.
op (Union[UnaryOps, BinaryOps, TernaryOps]) – The operation to be performed (UnaryOps, BinaryOps, TernaryOps).
srcs (LazyBuffer) – The input buffers for the operation.
arg (Optional[Any], optional) – An optional argument for certain operations, defaults to None.

Returns:

A new LazyBuffer with the result of the operation.

Return type:

LazyBuffer

srcs

The input buffers for the operation. Includes self.

Type:: LazyBuffer

out_device

The output device.

Type:: str

out_shape

The output shape.

Type:: Tuple[int, …]

out_dtype

The output data type.

Type:: DType

expand(arg: Tuple[Node | int, ...]) → LazyBuffer[source]

Expand the current LazyBuffer based on the given argument.

This function checks if the shape of the current LazyBuffer is equal to the provided argument. If so, it returns the LazyBuffer itself. If not, and if the LazyBuffer hasn’t been realized yet and its operation is an expansion operation, it returns the source of the operation with index 0, also expanded using the provided argument. Otherwise, it creates a new LazyBuffer by calling the movement_op function with the result of expanding the current state, the MovementOps.EXPAND operation, and the provided argument.

Parameters:

self (LazyBuffer) – The current LazyBuffer object.
arg (Tuple[sint, ...]) – A tuple containing sint objects representing the new shape of the LazyBuffer.

Returns:

A new LazyBuffer with the expanded shape.

Return type:

LazyBuffer

static fromCPU(x: ndarray) → LazyBuffer[source]

Create a new LazyBuffer object from a numpy array.

x

The numpy array to be used for creating the LazyBuffer.

Type:: np.ndarray

Returns:: A new LazyBuffer object created from the input numpy array.
Return type:: LazyBuffer

get_lazyops() → List[LazyOp][source]

Return an empty list of LazyOp objects.

This method is a placeholder and always returns an empty list. Subclasses may override this method to provide specific behavior.

Returns:: An empty list.
Return type:: List[LazyOp]

is_unrealized_const()[source]

Check whether this LazyBuffer is an unrealized constant.

Returns:: True if the buffer is unrealized and its base operation is a CONST LoadOp, False otherwise.

is_unrealized_contiguous_const()[source]

Check whether this LazyBuffer is an unrealized contiguous constant.

Returns:: True if the buffer is both unrealized and contiguous, False otherwise.

static loadop(op, shape: Tuple[Node | int, ...], dtype: DType, device: str, arg=None, src: LazyBuffer | None = None) → LazyBuffer[source]

Load operation factory method.

Creates and returns a new LazyBuffer object based on the given parameters. This is a static method and does not require an instance of the class.

op: The operation to be performed.

shape

The shape of the data to be loaded.

Type:: Tuple[sint, …]

dtype: The data type of the data to be loaded.

device

The device where the data will be loaded onto.

Type:: str

arg: An optional argument for the operation. Default is None.

src

An optional source LazyBuffer object. Default is None.

Type:: Optional[LazyBuffer]

Returns:: A new LazyBuffer object created with the given parameters.
Return type:: LazyBuffer

map_buffers(real_srcs: Mapping[Any, LazyBuffer | LazyOp])[source]

Retrieve the corresponding LazyBuffer or LazyOp object from the mapping of sources.

Parameters:: real_srcs (Mapping[Any, Union[LazyBuffer, LazyOp]]) – A mapping of objects, where keys are any hashable objects and values are either LazyBuffer or LazyOp.
Returns:: The corresponding LazyBuffer or LazyOp object. If the instance is not found in the mapping, it returns itself.
Return type:: Union[LazyBuffer, LazyOp]

pad(arg: Tuple[Tuple[int, int], ...]) → LazyBuffer[source]

Pad the buffer object.

This method pads the current buffer with specified start and end indices. If all padding values are 0, it simply returns the original buffer. If the buffer is not realized and the last operation was a pad operation, the pad operation is combined with the new one. Otherwise, a new movement operation is created with the pad argument.

Parameters:: arg (Tuple[Tuple[int, int], ...]) – A tuple of tuples containing start and end padding values.
Returns:: The padded buffer object.
Return type:: LazyBuffer

permute(arg: Tuple[int, ...]) → LazyBuffer[source]

Permute the current LazyBuffer based on the given argument.

This function is not yet implemented.

Parameters:

self (LazyBuffer) – The current LazyBuffer object.
arg (Tuple[int, ...]) – A tuple containing int objects representing the permutation of the dimensions.

Returns:

A new LazyBuffer with the permuted shape.

Return type:

LazyBuffer

r(op: ReduceOps, new_shape: Tuple[Node | int, ...]) → LazyBuffer[source]

Alias for _reduce_op.

This method is an alias for _reduce_op and provides another way to call it. It takes the same arguments as _reduce_op and returns a new LazyBuffer with reduced dimensions.

self

The LazyBuffer object on which the method is called.

Type:: LazyBuffer

op

The reduction operation to be applied.

Type:: ReduceOps

new_shape

The desired shape of the reduced LazyBuffer.

Type:: Tuple[sint, …]

Returns:: A new LazyBuffer object with reduced dimensions.
Return type:: LazyBuffer

property realized

Return whether this LazyBuffer is realized.

Returns:: True if the buffer is realized, False otherwise.

replace_with_movement_ops(ops: List[Tuple[MovementOps, Any]]) → LazyBuffer[source]

This method takes a list of tuples as an argument where each tuple contains a movement operation and its corresponding argument.

self

The lazy buffer instance on which the operations are to be performed.

Type:: LazyBuffer

ops

A list of tuples where each tuple contains a MovementOps enum member and its corresponding argument.

Type:: List[Tuple[MovementOps, Any]]

Returns:: The updated lazy buffer instance after performing all the movement operations in sequence.
Return type:: LazyBuffer

This method iterates over the list of tuples and for each tuple, it retrieves the corresponding function from the MOVEMENT_OPS_DISPATCHER dictionary using the MovementOps enum member as the key. The function is then called with the lazy buffer instance and its argument to perform the movement operation. The result of this operation is then stored back in the lazy buffer instance for subsequent operations. Once all the operations have been performed, the final updated lazy buffer instance is returned.

reshape(arg: Tuple[Node | int, ...]) → LazyBuffer[source]

Reshapes the buffer.

self

The current lazy buffer.

Type:: LazyBuffer

arg

The new shape for the buffer.

Type:: Tuple[sint, …]

Returns:: The reshaped lazy buffer.
Return type:: LazyBuffer

schedule(seen: Set[LazyBuffer] | None = None) → List[ScheduleItem][source]

Schedules the computation of this lazy buffer.

Parameters:: seen (Optional[Set[LazyBuffer]]) – Set of already scheduled buffers. Defaults to None.
Returns:: A list of schedule items for this buffer’s computations.
Return type:: List[ScheduleItem]

seen

Set of already scheduled buffers. Defaults to an empty set if not provided.

Type:: Optional[Set[LazyBuffer]]

ret

List of schedule items for this buffer’s computations.

Type:: List[ScheduleItem]

var_vals

Merged dictionary of variable values from this buffer and its operand buffers.

Type:: Dict

op

The abstract syntax tree node representing the operation to be performed on this buffer.

Type:: ASTNode

base_bufs

List of base buffers for this buffer’s computation.

Type:: List[Buffer]

shrink(arg: Tuple[Tuple[Node | int, Node | int], ...]) → LazyBuffer[source]

Shrinks the buffer based on the provided arguments.

self

The lazy buffer to be shrunk.

Type:: LazyBuffer

arg

A tuple of tuples containing the start and end indices for each dimension of the buffer.

Type:: Tuple[Tuple[sint, sint], …]

Returns:: The shrunken lazy buffer.
Return type:: LazyBuffer

stride(arg: Tuple[int, ...]) → LazyBuffer[source]

Applies a stride operation to the buffer based on the provided arguments.

self

The lazy buffer to have the stride operation applied to.

Type:: LazyBuffer

arg

A tuple of integers representing the strides for each dimension of the buffer.

Type:: Tuple[int, …]

Returns:: The lazy buffer after the stride operation has been applied.
Return type:: LazyBuffer

tinygrad.lazy.MOVEMENT_OPS_DISPATCHER: Dict[MovementOps, Callable] = {MovementOps.EXPAND: <function LazyBuffer.expand>, MovementOps.PAD: <function LazyBuffer.pad>, MovementOps.PERMUTE: <function LazyBuffer.permute>, MovementOps.RESHAPE: <function LazyBuffer.reshape>, MovementOps.SHRINK: <function LazyBuffer.shrink>, MovementOps.STRIDE: <function LazyBuffer.stride>}

This dictionary acts as a dispatcher for various movement operations on LazyBuffer objects. It maps each operation to its corresponding function in the LazyBuffer class.

tinygrad.lazy.MovementOps[source]

An enumeration of different movement operations like reshape, expand, shrink, permute, pad, and stride.

Type:: Enum

tinygrad.lazy.LazyBuffer[source]

The LazyBuffer class which contains the methods corresponding to each operation in this dispatcher.

Type:: Class

tinygrad.lazy.Dict[MovementOps, Callable]

Type hint indicating that keys are movement operations and values are their corresponding functions in LazyBuffer.

Type:: Type Hinting

The dictionary keys

MovementOps.RESHAPE: Corresponds to the reshape operation, mapped to LazyBuffer.reshape method.

MovementOps.EXPAND: Corresponds to the expand operation, mapped to LazyBuffer.expand method.

MovementOps.SHRINK: Corresponds to the shrink operation, mapped to LazyBuffer.shrink method.

MovementOps.PERMUTE: Corresponds to the permute operation, mapped to LazyBuffer.permute method.

MovementOps.PAD: Corresponds to the pad operation, mapped to LazyBuffer.pad method.

MovementOps.STRIDE: Corresponds to the stride operation, mapped to LazyBuffer.stride method.

Create a lazy buffer for the given device, shape tracker, operation type, operation, data type, and optional base.

Parameters:

device (str) – The device to create the lazy buffer on.
st (ShapeTracker) – The shape tracker to use.
optype (OpType) – The operation type.
op (LazyOp) – The lazy operation.
dtype (DType) – The data type of the lazy buffer.
base (Optional[LazyBuffer]) – The optional base for the lazy buffer. Default is None.

Returns:

The created lazy buffer.

Return type:

LazyBuffer

tinygrad.lazy.get_movementroot(root: LazyBuffer, allow_contiguous=False) → LazyBuffer[source]

Recursively retrieve the root of a movement operation or contiguous data.

This function is used to locate the origin of a series of operations by traversing back through the chain of operations until it reaches the original source buffer. It will continue this process as long as the current root is not realized and its operation type matches specific criteria.

Parameters:

root (LazyBuffer) – The current root node in the operation tree.
allow_contiguous (bool, optional) – A flag indicating whether to include operations that result in contiguous data, defaults to False.

Returns:

The original root node of the operation tree or the current root if it is realized or does not meet the specified criteria.

Return type:

LazyBuffer

tinygrad.lazy.get_movementroot_contiguous(x: LazyBuffer) → LazyBuffer[source]

Recursively obtain the root of movement for a contiguous operation.

This function is used to identify and return the root of a series of operations that lead up to a contiguous operation on a lazy buffer. It does this by checking if the current operation is a contiguous operation, and if it is not, recursively calling itself on the source operation until it finds the root of the contiguous operation.

Parameters:: x (LazyBuffer) – The input lazy buffer to check for the root of its movement.
Returns:: The root of the movement for the contiguous operation.
Return type:: LazyBuffer

tinygrad.lazy.vars_from_ast(ast: LazyOp) → List[Variable][source]

Retrieve variables from abstract syntax tree (AST).

This function extracts all unique variables from the AST by leveraging the LazyOp operations that belong to BufferOps. The resulting set of variables is then sorted based on their expression string representation.

tinygrad.lazy.ast

Abstract syntax tree object to extract variables from.

Type:: LazyOp

Returns:: Sorted list of unique variables extracted from AST.
Return type:: List[Variable]