tinygrad mlops
Note
You likely want the upstream tinygrad, not tinygrab. Tinygrab contains AI generated docstrings for a tinygrad snapshot. Upstream: https://tinygrad.org
- class tinygrad.mlops.Add(device: str, *tensors: Tensor)[source]
Bases:
Function
This class represents an addition operation. It takes two arguments, x and y, which are LazyBuffer objects. The forward method adds these two inputs using BinaryOps.ADD operator and returns the result as a LazyBuffer object. The backward method calculates gradients with respect to the input tensors.
- x
The first input buffer.
- Type:
LazyBuffer
- y
The second input buffer.
- Type:
LazyBuffer
- backward(grad_output: LazyBuffer) Tuple[LazyBuffer | None, LazyBuffer | None] [source]
This method calculates the gradient of the output with respect to the input tensors.
- grad_output
The gradient output buffer.
- Type:
LazyBuffer
- Returns:
A tuple containing the gradients with respect to the first and second input tensors, if they are required.
- Return type:
Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]
- class tinygrad.mlops.Cast(device: str, *tensors: Tensor)[source]
Bases:
Function
This class defines a function for casting tensors to a specified data type.
- forward[source]
This method accepts a LazyBuffer, a DType and a boolean flag bitcast. It returns the casted tensor.
- Type:
method
- backward[source]
This method accepts a LazyBuffer and returns its casted version with respect to the input dtype and bitcast.
- Type:
method
- backward(grad_output: LazyBuffer) LazyBuffer [source]
Casts the gradient tensor to the input data type and performs a bitcast operation if necessary.
- Parameters:
grad_output (LazyBuffer) – The gradient tensor.
- Returns:
The casted version of the gradient tensor with respect to the original input tensor data type and bitcast.
- Return type:
LazyBuffer
- forward(x: LazyBuffer, dtype: DType, bitcast: bool = False) LazyBuffer [source]
Casts the input tensor to a specified data type.
- Parameters:
x (LazyBuffer) – The input tensor.
dtype (DType) – The target data type.
bitcast (bool) – Whether to perform a bitcast operation. Default is False.
- Returns:
The casted version of the input tensor.
- Return type:
LazyBuffer
- class tinygrad.mlops.Contiguous(device: str, *tensors: Tensor)[source]
Bases:
Function
This class defines a function for ensuring tensor is contiguous in memory.
- forward[source]
This method accepts a LazyBuffer and returns a LazyBuffer that is contiguous.
- Type:
method
- backward[source]
This method accepts a LazyBuffer and returns it without any modification.
- Type:
method
- class tinygrad.mlops.ContiguousBackward(device: str, *tensors: Tensor)[source]
Bases:
Function
This class defines a function for ensuring tensor is contiguous in memory during backpropagation.
- forward[source]
This method accepts a LazyBuffer and returns it without any modification.
- Type:
method
- class tinygrad.mlops.Div(device: str, *tensors: Tensor)[source]
Bases:
Function
The Div class represents a division function. It has methods for forward and backward propagation.
- x
The input buffer representing the numerator.
- Type:
LazyBuffer
- y
The input buffer representing the denominator.
- Type:
LazyBuffer
- backward(grad_output: LazyBuffer) Tuple[LazyBuffer | None, LazyBuffer | None] [source]
Perform backward propagation for the division operation.
- Parameters:
grad_output (LazyBuffer) – The gradient of the output buffer.
- Returns:
A tuple containing the gradients with respect to the numerator and denominator. If an input does not require a gradient, None is returned for that position.
- Return type:
Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]
- forward(x: LazyBuffer, y: LazyBuffer) LazyBuffer [source]
Perform forward propagation for the division operation.
- Parameters:
x (LazyBuffer) – The input buffer representing the numerator.
y (LazyBuffer) – The input buffer representing the denominator.
- Returns:
The result of the division operation.
- Return type:
LazyBuffer
- class tinygrad.mlops.Exp(device: str, *tensors: Tensor)[source]
Bases:
Function
This class represents an exponential function. It has methods for forward and backward operations.
- backward(grad_output: LazyBuffer) LazyBuffer [source]
Backward method: Performs the backward operation to calculate gradient of the exponential function.
- Parameters:
grad_output (LazyBuffer) – Gradient buffer from previous operation
- Returns:
Output buffer after applying the backward exponential operation
- Return type:
LazyBuffer
- class tinygrad.mlops.Expand(device: str, *tensors: Tensor)[source]
Bases:
Function
Class for expanding a tensor to a given shape.
- input_shape
The original shape of the input tensor before expansion.
- Type:
Tuple[int, …]
- class tinygrad.mlops.Flip(device: str, *tensors: Tensor)[source]
Bases:
Function
The Flip class, a subclass of Function. This class defines two methods: forward and backward.
- arg
A tuple containing the stride arguments for flipping the input tensor. It is calculated by iterating over the length of the input tensor’s shape and negating 1 if the current index is in the axis parameter, otherwise positively keeping 1.
- Type:
tuple
- backward(grad_output: LazyBuffer) LazyBuffer [source]
Flip the gradient tensor along specified axes, mirroring the forward operation.
- Parameters:
grad_output (LazyBuffer) – The gradient tensor to be flipped.
- Returns:
The flipped gradient tensor.
- Return type:
LazyBuffer
- forward(x: LazyBuffer, axis: Tuple[int, ...]) LazyBuffer [source]
Flip the input tensor along specified axes.
- Parameters:
x (LazyBuffer) – The input tensor to be flipped.
axis (Tuple[int, ...]) – A tuple of integers representing the axes along which to flip the tensor.
- Returns:
The flipped output tensor.
- Return type:
LazyBuffer
- class tinygrad.mlops.Less(device: str, *tensors: Tensor)[source]
Bases:
Function
This class represents a less than operation. It takes two arguments, x and y, which are LazyBuffer objects. The forward method compares these two inputs using BinaryOps.CMPLT operator and returns the result as a LazyBuffer object.
- x
The first input buffer.
- Type:
LazyBuffer
- y
The second input buffer.
- Type:
LazyBuffer
- class tinygrad.mlops.Log(device: str, *tensors: Tensor)[source]
Bases:
Function
Function for calculating the logarithm of a given input.
- x
Input buffer on which operations are performed.
- Type:
LazyBuffer
- class tinygrad.mlops.Max(device: str, *tensors: Tensor)[source]
Bases:
Function
Class for the Max Function.
- x
The input tensor.
- Type:
LazyBuffer
- ret
The result of applying the max function to x.
- Type:
LazyBuffer
- class tinygrad.mlops.Mul(device: str, *tensors: Tensor)[source]
Bases:
Function
Function class for element-wise multiplication operation.
- x
First input tensor.
- Type:
LazyBuffer
- y
Second input tensor.
- Type:
LazyBuffer
- backward(grad_output: LazyBuffer) Tuple[LazyBuffer | None, LazyBuffer | None] [source]
Backward pass of the element-wise multiplication operation.
- Parameters:
grad_output (LazyBuffer) – Gradient tensor from previous operation.
- Returns:
- Gradients with respect to input tensors x and y.
If the gradient is not needed for a certain input tensor, None is returned instead.
- Return type:
Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]
- class tinygrad.mlops.Neg(device: str, *tensors: Tensor)[source]
Bases:
Function
Neg class for negating a given value.
- forward[source]
Method to perform the forward pass of the negation operation on the input x.
- Type:
method
- backward[source]
Method to perform the backward pass of the negation operation, used during backpropagation.
- Type:
method
- class tinygrad.mlops.Pad(device: str, *tensors: Tensor)[source]
Bases:
Function
The Pad class is used to add padding to a given tensor. It supports forward and backward propagation.
- narg
A tuple storing the new arguments for padding.
- Type:
tuple
- backward(grad_output: LazyBuffer) LazyBuffer [source]
The backward method removes the padding from the gradient tensor.
- Parameters:
grad_output (LazyBuffer) – The gradient tensor with padding.
- Returns:
The gradient tensor after removing padding.
- Return type:
LazyBuffer
- forward(x: LazyBuffer, arg: Tuple[Tuple[int, int], ...]) LazyBuffer [source]
The forward method applies padding to the input tensor.
- Parameters:
x (LazyBuffer) – The input tensor.
arg (tuple of tuples) – A tuple containing two-element tuples that specify the padding for each dimension.
- Returns:
The output tensor after padding.
- Return type:
LazyBuffer
- class tinygrad.mlops.Permute(device: str, *tensors: Tensor)[source]
Bases:
Function
The Permute class is a subclass of the Function class, used for permutation operations on tensors.
- input_order
The order of the elements in the original tensor. This attribute is set during the forward pass.
- Type:
Tuple[int, …]
- backward(grad_output: LazyBuffer) LazyBuffer [source]
Permute the gradient tensor according to the inverse of the input order.
- Parameters:
grad_output (LazyBuffer) – The gradient tensor to be permuted.
- Returns:
The permuted gradient tensor.
- Return type:
LazyBuffer
- forward(x: LazyBuffer, order: Tuple[int, ...]) LazyBuffer [source]
Permute the input tensor according to a specified order.
- Parameters:
x (LazyBuffer) – The input tensor to be permuted.
order (Tuple[int, ...]) – The order of the elements in the desired output tensor.
- Returns:
The output tensor after permutation.
- Return type:
LazyBuffer
- class tinygrad.mlops.Relu(device: str, *tensors: Tensor)[source]
Bases:
Function
The Relu class implements the Rectified Linear Unit (ReLU) activation function.
- x
Input tensor to which ReLU is applied.
- Type:
LazyBuffer
- ret
Output tensor after applying ReLU.
- Type:
LazyBuffer
- class tinygrad.mlops.Reshape(device: str, *tensors: Tensor)[source]
Bases:
Function
This class represents a function for reshaping tensors.
- input_shape
The shape of the input tensor before reshaping.
- Type:
Tuple[int, …]
- backward(grad_output: LazyBuffer) LazyBuffer [source]
Reshape the gradient tensor during the backward pass.
- Parameters:
grad_output (LazyBuffer) – The gradient tensor to be reshaped.
- Returns:
The reshaped gradient tensor with the same shape as the input tensor from the forward pass.
- Return type:
LazyBuffer
- forward(x: LazyBuffer, shape: Tuple[int, ...]) LazyBuffer [source]
Reshape the input tensor during the forward pass.
- Parameters:
x (LazyBuffer) – The input tensor to be reshaped.
shape (Tuple[int, ...]) – The desired shape of the output tensor.
- Returns:
The reshaped output tensor.
- Return type:
LazyBuffer
- class tinygrad.mlops.Shrink(device: str, *tensors: Tensor)[source]
Bases:
Function
The Shrink class implements the forward and backward methods for a shrinking function.
- narg
A tuple storing the result of a zip operation between x.shape and arg, where each element is a 2-tuple containing integers.
- Type:
tuple
- backward(grad_output: LazyBuffer) LazyBuffer [source]
The backward method takes a lazy buffer grad_output and applies the pad operation to it using the class attribute narg. It first checks if all elements in narg are integers. If not, it raises an assertion error stating that symbolic shrink does not support backward.
- Parameters:
grad_output (LazyBuffer) – The input lazy buffer for the backward operation.
- Returns:
The result of calling the pad method on grad_output with self.narg as an argument.
- Return type:
LazyBuffer
- Raises:
AssertionError – If not all elements in narg are integers.
- forward(x: LazyBuffer, arg: Tuple[Tuple[Node | int, Node | int], ...]) LazyBuffer [source]
The forward method takes a lazy buffer x and a tuple of tuples arg, where each inner tuple contains two integers. It computes a new tuple by applying the shrinking operation to the elements of x.shape. This result is stored in the class attribute narg. Finally, it returns the result of calling the shrink method on x with arg as an argument.
- Parameters:
x (LazyBuffer) – The input lazy buffer.
arg (Tuple[Tuple[sint, sint], ...]) – A tuple containing 2-tuples of integers.
- Returns:
The result of calling the shrink method on x with arg as an argument.
- Return type:
LazyBuffer
- class tinygrad.mlops.Sigmoid(device: str, *tensors: Tensor)[source]
Bases:
Function
Sigmoid Activation Function.
This class implements the sigmoid activation function and its derivative. The forward method computes the sigmoid of the input tensor x. The backward method, on the other hand, computes the gradient of the loss function with respect to the input of the sigmoid function.
- ret
The result of the forward pass computation.
- Type:
LazyBuffer
- backward(grad_output: LazyBuffer) LazyBuffer [source]
Backward method for computing the gradient of the loss function.
This method computes the gradient of the loss function with respect to the input of the sigmoid function. It uses the following mathematical formula to compute the result:
grad_input = sigmoid(x) * (1 - sigmoid(x)) * grad_output
- Parameters:
grad_output (LazyBuffer) – The gradient of the loss function with respect to the output.
- Returns:
The gradient of the loss function with respect to the input.
- Return type:
LazyBuffer
- forward(x: LazyBuffer) LazyBuffer [source]
Forward method for computing sigmoid activation.
This method computes the sigmoid function of input tensor x. It uses the following mathematical formula to compute the result:
sigmoid(x) = 1 / (1 + exp(-x * (-1/log(2))))
- Parameters:
x (LazyBuffer) – The input tensor.
- Returns:
The output tensor after applying the sigmoid function.
- Return type:
LazyBuffer
- class tinygrad.mlops.Sin(device: str, *tensors: Tensor)[source]
Bases:
Function
Implementation of the sine function in radians.
- x
Input buffer for which the sine operation is applied.
- Type:
LazyBuffer
- backward(grad: LazyBuffer) LazyBuffer [source]
Apply the gradient of the sine function to input LazyBuffer and return result.
The gradient of sin(x) is cos(x). Hence, we apply a cos operation here.
- Parameters:
grad (LazyBuffer) – Gradient buffer for which the operation is applied.
- Returns:
Output buffer after applying the gradient operation.
- Return type:
LazyBuffer
- class tinygrad.mlops.Sqrt(device: str, *tensors: Tensor)[source]
Bases:
Function
Function for calculating the square root of a given LazyBuffer.
- ret
The result of the forward operation, which is the input value’s square root.
- Type:
LazyBuffer
- backward(grad_output: LazyBuffer) LazyBuffer [source]
Calculate the gradient of the square root function with respect to its input.
- Parameters:
grad_output (LazyBuffer) – The gradient output from previous operations.
- Returns:
The gradient with respect to the input of the square root function, calculated as grad_output / (2 * sqrt(x)).
- Return type:
LazyBuffer
- forward(x: LazyBuffer) LazyBuffer [source]
Calculate the square root of the input LazyBuffer and store it in ret.
- Parameters:
x (LazyBuffer) – The input value to calculate the square root for.
- Returns:
The result of the forward operation, which is the input value’s square root.
- Return type:
LazyBuffer
- class tinygrad.mlops.Sub(device: str, *tensors: Tensor)[source]
Bases:
Function
This class represents a subtraction operation. It takes two arguments, x and y, which are LazyBuffer objects. The forward method subtracts the second input from the first using BinaryOps.SUB operator and returns the result as a LazyBuffer object. The backward method calculates gradients with respect to the input tensors.
- x
The first input buffer.
- Type:
LazyBuffer
- y
The second input buffer.
- Type:
LazyBuffer
- backward(grad_output: LazyBuffer) Tuple[LazyBuffer | None, LazyBuffer | None] [source]
This method calculates the gradient of the output with respect to the input tensors.
- grad_output
The gradient output buffer.
- Type:
LazyBuffer
- Returns:
A tuple containing the gradients with respect to the first and second input tensors, if they are required.
- Return type:
Tuple[Optional[LazyBuffer], Optional[LazyBuffer]]
- class tinygrad.mlops.Sum(device: str, *tensors: Tensor)[source]
Bases:
Function
Class for summing elements of a LazyBuffer and backward operation to expand gradients.
- input_shape
The shape of the original LazyBuffer used in the forward method.
- Type:
Tuple[int, …]
- backward(grad_output: LazyBuffer) LazyBuffer [source]
Backward method to expand the gradient (grad_output) according to the input shape of the forward method.
- Parameters:
grad_output (LazyBuffer) – The gradient from previous operation.
- Returns:
The expanded gradient according to the input shape in the forward method.
- Return type:
LazyBuffer
- forward(x: LazyBuffer, new_shape: Tuple[int, ...]) LazyBuffer [source]
Forward method to sum elements of a LazyBuffer based on the new shape provided.
- Parameters:
x (LazyBuffer) – The input LazyBuffer.
new_shape (Tuple[int, ...]) – The desired output shape after the summation operation.
- Returns:
The result of the summation operation with the specified new shape.
- Return type:
LazyBuffer
- class tinygrad.mlops.Where(device: str, *tensors: Tensor)[source]
Bases:
Function
The Where class is used to perform element-wise conditional operations.
- x
The input buffer from which the operation will be performed.
- Type:
LazyBuffer
- backward(grad_output: LazyBuffer) Tuple[None, LazyBuffer | None, LazyBuffer | None] [source]
Backward method for the Where class.
This method performs the backward pass of the element-wise conditional operation.
- Parameters:
grad_output (LazyBuffer) – The gradient output used in the backward pass calculation.
- Returns:
A tuple containing None as the first element and optional LazyBuffer objects for the second and third operands, depending on whether their gradients are needed or not.
- Return type:
Tuple[None, Optional[LazyBuffer], Optional[LazyBuffer]]
- forward(x: LazyBuffer, y: LazyBuffer, z: LazyBuffer) LazyBuffer [source]
Forward method for the Where class.
This method performs the element-wise conditional operation during forward propagation.
- Parameters:
x (LazyBuffer) – The input buffer from which the operation will be performed.
y (LazyBuffer) – The first operand used in the operation.
z (LazyBuffer) – The second operand used in the operation.
- Returns:
The result of the element-wise conditional operation.
- Return type:
LazyBuffer
- class tinygrad.mlops.Zero(device: str, *tensors: Tensor)[source]
Bases:
Function
Zero class for performing basic operations related to zero.
This class inherits from the Function class and overrides its methods.
- backward(grad: LazyBuffer) LazyBuffer [source]
Backward method for applying the backward pass of the operation.
This method takes a lazy buffer as input and returns another lazy buffer with all elements set to zero.
- Parameters:
grad (LazyBuffer) – The gradient lazy buffer from previous operations.
- Returns:
The output lazy buffer with all elements set to zero.
- Return type:
LazyBuffer
- forward(x: LazyBuffer) LazyBuffer [source]
Forward method for applying the forward pass of the operation.
This method takes a lazy buffer as input and returns another lazy buffer with all elements set to zero.
- Parameters:
x (LazyBuffer) – The input lazy buffer.
- Returns:
The output lazy buffer with all elements set to zero.
- Return type:
LazyBuffer