tinygrad ops

Note

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

class tinygrad.ops.BinaryOps(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: Enum

This class defines an enumeration for various binary operations.

The enumerations are:

ADD - Addition operation

SUB - Subtraction operation

MUL - Multiplication operation

DIV - Division operation

MAX - Maximum of two operands operation

MOD - Modulo operation

CMPLT - Comparison less than operation

ADD = 1

CMPLT = 7

DIV = 4

MAX = 5

MOD = 6

MUL = 3

SUB = 2

class tinygrad.ops.BufferOps(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: Enum

This class represents the different types of buffer operations.

LOAD: A buffer operation that loads a value into a buffer.

CONST: A buffer operation that sets a constant value in a buffer.

STORE: A buffer operation that stores a value from a buffer.

CONST = 2

LOAD = 1

STORE = 3

class tinygrad.ops.ConstBuffer(val: 'Union[int, float]', dtype: 'DType', st: 'ShapeTracker')[source]

Bases: object

dtype: DType

st: ShapeTracker

val: int | float

class tinygrad.ops.FlopCounter(shape: Tuple[int, ...], dtype: DType, flops: int, mem: Dict[int, int])[source]

Bases: object

Data class for tracking Floating Point Operations (FLOPs) and memory usage.

shape

Tuple[int, …] The shape of the data structure being tracked.

Type:: Tuple[int, …]

dtype

DType The data type of the elements in the data structure.

Type:: tinygrad.helpers.DType

flops

int The number of floating point operations performed.

Type:: int

mem

Dict[int, int] A dictionary mapping memory location IDs to their associated memory usage.

Type:: Dict[int, int]

consume_flops()[source]

Consume all tracked floating point operations and return their count.

Postconditions:: The flops attribute is set to 0 after this function call.

Returns:: The number of consumed floating point operations.
Return type:: int

dtype: DType

flops: int

mem: Dict[int, int]

property mem_estimate

Calculate the total memory usage of the data structure.

Returns:: The sum of all memory usages in the mem attribute.

shape: Tuple[int, ...]

class tinygrad.ops.LazyOp(op: Op, src: Tuple[LazyOp | LazyBuffer, ...], arg: Any = None)[source]

Bases: object

Data class for lazy operations.

op

Operation to be performed.

Type:: Op

src

Source data for the operation.

Type:: Tuple[Union[LazyOp, LazyBuffer], …]

arg

Optional argument for the operation. Default is None.

Type:: Any

Properties:: hash (int): Cached property for the hash of the object. buffers (Tuple[LazyBuffer, …]): Cached property for the buffers in the source data.

map_buffers(real_srcs: Mapping[Any, Union[LazyBuffer, LazyOp]]) -> LazyOp: Maps the buffers to real sources.

get_lazyops() → List[LazyOp][source]: Returns a list of lazy operations.

replace_with_movement_ops(ops: List[Tuple[MovementOps, Tuple[Any, …]]]) -> “LazyBuffer”: Replaces the operation with movement operations.

arg: Any = None

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

Cached property for the buffers in the source data.

Returns:: The buffers in the source data.
Return type:: Tuple[LazyBuffer, …]

property children

Get the children attribute of self.

Returns:: The operation is not implemented yet.
Return type:: NotImplementedError
Raises:: NotImplementedError – This method is not implemented at the moment.

expand(_)[source]

Expands self based on input arguments.

Parameters:: _ – The argument(s) to be used for expansion.
Returns:: The operation is not implemented yet.
Return type:: NotImplementedError
Raises:: NotImplementedError – This method is not implemented at the moment.

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

Returns a list of lazy operations.

Returns:: A list of lazy operations.
Return type:: List[LazyOp]

property hash

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

Maps the buffers to real sources.

Parameters:: real_srcs (Mapping[Any, Union[LazyBuffer, LazyOp]]) – The mapping of buffers to real sources.
Returns:: A new lazy operation with mapped buffers.
Return type:: LazyOp

op: Op

pad(_)[source]

Pads self with specified values based on input arguments.

Parameters:: _ – The argument(s) to be used for padding.
Returns:: The operation is not implemented yet.
Return type:: NotImplementedError
Raises:: NotImplementedError – This method is not implemented at the moment.

permute(_)[source]

Permutes self based on input arguments.

Parameters:: _ – The argument(s) to be used for permutation.
Returns:: The operation is not implemented yet.
Return type:: NotImplementedError
Raises:: NotImplementedError – This method is not implemented at the moment.

property realized

Get the realized attribute of self.

Returns:: The operation is not implemented yet.
Return type:: NotImplementedError
Raises:: NotImplementedError – This method is not implemented at the moment.

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

Replaces the operation with movement operations.

Parameters:: ops (List[Tuple[MovementOps, Tuple[Any, ...]]]) – List of tuples containing movement operations and their arguments.
Returns:: A new lazy buffer with replaced operations.
Return type:: LazyBuffer

reshape(_)[source]

Reshapes self based on input arguments.

Parameters:: _ – The argument(s) to be used for reshaping.
Returns:: The operation is not implemented yet.
Return type:: NotImplementedError
Raises:: NotImplementedError – This method is not implemented at the moment.

shrink(_)[source]

Shrinks self based on input arguments.

Parameters:: _ – The argument(s) to be used for shrinking.
Returns:: The operation is not implemented yet.
Return type:: NotImplementedError
Raises:: NotImplementedError – This method is not implemented at the moment.

src: Tuple[LazyOp | LazyBuffer, ...]

property st

Retrieve the st attribute of an instance of a class.

Returns:: The operation is not implemented yet.
Return type:: NotImplementedError
Raises:: NotImplementedError – This method is not implemented at the moment.

stride(_)[source]

Applies stride to self based on input arguments.

Parameters:: _ – The argument(s) to be used for applying stride.
Returns:: The operation is not implemented yet.
Return type:: NotImplementedError
Raises:: NotImplementedError – This method is not implemented at the moment.

class tinygrad.ops.LoadOps(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: Enum

Enumeration for load operation types.

EMPTY

Represents an empty load operation.

Type:: auto()

CONST

Represents a constant load operation, where the value is known and fixed.

Type:: auto()

FROM

Represents a load operation from another source.

Type:: auto()

CONTIGUOUS

Represents a contiguous load operation, where elements are stored in memory sequentially.

Type:: auto()

CUSTOM

Represents a custom load operation, where the behavior is user-defined or specified.

Type:: auto()

CONST = 2

CONTIGUOUS = 4

CUSTOM = 5

EMPTY = 1

FROM = 3

class tinygrad.ops.MemBuffer(idx: 'int', dtype: 'DType', st: 'ShapeTracker')[source]

Bases: object

dtype: DType

idx: int

st: ShapeTracker

class tinygrad.ops.MovementOps(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: Enum

Enum class for various movement operations.

RESHAPE: Represents a reshaping operation.

PERMUTE: Represents a permuting operation.

EXPAND: Represents an expanding operation.

PAD: Represents a padding operation.

SHRINK: Represents a shrinking operation.

STRIDE: Represents a striding operation.

AS_STRIDED: Represents an as-strided operation.

AS_STRIDED = 7

EXPAND = 3

PAD = 4

PERMUTE = 2

RESHAPE = 1

SHRINK = 5

STRIDE = 6

tinygrad.ops.Op

This module defines the operations that can be performed on tensors.

tinygrad.ops.Op

The union of all operation types supported by this module.

Type:: Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, LoadOps, TernaryOps, BufferOps]

tinygrad.ops.OpType

The union of all operation types as class objects.

Type:: Union[Type[UnaryOps], Type[BinaryOps], Type[ReduceOps], Type[MovementOps], Type[LoadOps], Type[TernaryOps], Type[BufferOps]]

alias of Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, LoadOps, TernaryOps, BufferOps]

class tinygrad.ops.ReduceOps(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: Enum

Define an enumeration for reduce operations.

The ReduceOps class is an enumeration that defines the possible reduce operations available for use in other parts of a program. It includes two primary options - SUM and MAX. These are defined using the built-in Enum class from Python’s standard library.

MAX = 2

SUM = 1

class tinygrad.ops.ScheduleItem(ast: 'LazyOp', out: 'LazyBuffer', inputs: 'Tuple[LazyBuffer, ...]', var_vals: 'Dict[Variable, int]')[source]

Bases: object

ast: LazyOp

inputs: Tuple[LazyBuffer, ...]

out: LazyBuffer

var_vals: Dict[Variable, int]

class tinygrad.ops.TernaryOps(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: Enum

This is an Enum class for representing ternary operations.

MULACC: Represents a multiplication accumulation operation.

WHERE: Represents a where operation.

MULACC = 1

WHERE = 2

class tinygrad.ops.UnaryOps(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: Enum

The EXP2 operation represents the base-2 exponential function. This operation calculates 2 raised to the power of the input value.

The LOG2 operation represents the base-2 logarithm function. This operation calculates the logarithm base 2 of the input value.

The CAST operation is used for type conversion or casting. For example, this operation can be used to convert a float number to an integer.

The SIN operation calculates the sine of the input value. It represents the standard mathematical sin function.

The SQRT operation calculates the square root of the input value. This is a standard mathematical operation that finds the number that, when multiplied by itself, gives the original input value.

The RECIP operation calculates the reciprocal of the input value. This means it calculates the number needed to multiply the input in order to get 1 as a result (e.g., if the input is 2, the output would be 0.5).

Note: In rdna3, only RECIP is available and not DIV or POW.

The NEG operation calculates the negation of the input value. This means it changes the sign of the input number (e.g., if the input is 5, the output would be -5).

CAST = 3

EXP2 = 1

LOG2 = 2

NEG = 7

RECIP = 6

SIN = 4

SQRT = 5

tinygrad.ops.get_lazyop_info(ast: LazyOp) → FlopCounter[source]

Run the AST using an InterpretedFlopCounter and return the result.

tinygrad.ops.ast

The abstract syntax tree to be evaluated.

Type:: LazyOp

Returns:: The result of running the AST through the InterpretedFlopCounter.
Return type:: FlopCounter