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
- 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:
- 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
- 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
- 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.
- 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.
- 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.
- 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.
- 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
- 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
- inputs: Tuple[LazyBuffer, ...]
- out: LazyBuffer
- 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.
- Returns:
The result of running the AST through the InterpretedFlopCounter.
- Return type: