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]]

forward(x: LazyBuffer, y: LazyBuffer) → LazyBuffer[source]

Forward method to be implemented by subclasses. This method is called when the function is called in forward mode.

Raises:: NotImplementedError – If not overridden in a subclass.

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

backward(grad_output: LazyBuffer) → LazyBuffer[source]

Returns the gradient output without any modification.

Parameters:: grad_output (LazyBuffer) – The gradient tensor.
Returns:: The same as the input grad_output tensor.
Return type:: LazyBuffer

forward(x: LazyBuffer) → LazyBuffer[source]

Returns a contiguous version of the input tensor.

Parameters:: x (LazyBuffer) – The input tensor.
Returns:: A contiguous version of the input tensor.
Return type:: LazyBuffer

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

backward[source]

This method accepts a LazyBuffer and returns its contiguous version.

Type:: method

backward(grad_output: LazyBuffer) → LazyBuffer[source]

Returns a contiguous version of the gradient tensor.

Parameters:: grad_output (LazyBuffer) – The gradient tensor.
Returns:: A contiguous version of the input grad_output tensor.
Return type:: LazyBuffer

forward(x: LazyBuffer) → LazyBuffer[source]

Returns the input tensor without any modifications.

Parameters:: x (LazyBuffer) – The input tensor.
Returns:: The same as the input tensor.
Return type:: LazyBuffer

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

forward(x: LazyBuffer) → LazyBuffer[source]

Forward method: Performs the exponential operation on input x.

Parameters:: x (LazyBuffer) – Input buffer
Returns:: Output buffer after applying the exponential function
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, …]

forward()[source]: Expands a LazyBuffer to a specified shape.

backward()[source]: Reduces the gradient back to the original input shape using a SUM operation.

backward(grad_output: LazyBuffer) → LazyBuffer[source]

Reduces the gradient back to the original input shape using a SUM operation.

Parameters:: grad_output (LazyBuffer) – The gradient tensor to be reduced.
Returns:: The reduced gradient tensor.
Return type:: LazyBuffer

forward(x: LazyBuffer, shape: Tuple[int, ...]) → LazyBuffer[source]

Expands a LazyBuffer to a specified shape.

Parameters:

x (LazyBuffer) – The input tensor to be expanded.
shape (Tuple[int, ...]) – The target shape for the expansion.

Returns:

The expanded tensor.

Return type:

LazyBuffer

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

forward(x: LazyBuffer, y: LazyBuffer) → LazyBuffer[source]

Forward method to be implemented by subclasses. This method is called when the function is called in forward mode.

Raises:: NotImplementedError – If not overridden in a subclass.

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

backward(grad_output: LazyBuffer) → LazyBuffer[source]

Backward method to compute the gradient of the logarithm function.

Parameters:: grad_output (LazyBuffer) – Gradient output buffer.
Returns:: Gradient buffer after performing operations.
Return type:: LazyBuffer

forward(x: LazyBuffer) → LazyBuffer[source]

Forward method to compute the natural logarithm of the input.

Parameters:: x (LazyBuffer) – Input buffer.
Returns:: Output buffer after performing operations.
Return 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

backward(grad_output: LazyBuffer) → LazyBuffer[source]

Backward pass for the Max Function.

Parameters:: grad_output (LazyBuffer) – The gradient output tensor.
Returns:: Gradient input tensor.
Return type:: LazyBuffer

forward(x: LazyBuffer, new_shape: Tuple[int, ...]) → LazyBuffer[source]

Forward pass for the Max Function.

Parameters:

x (LazyBuffer) – Input tensor.
new_shape (Tuple[int, ...]) – The desired shape of the output.

Returns:

Result of applying the max function to x.

Return 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]]

forward(x: LazyBuffer, y: LazyBuffer) → LazyBuffer[source]

Forward pass of the element-wise multiplication operation.

Parameters:

x (LazyBuffer) – First input tensor.
y (LazyBuffer) – Second input tensor.

Returns:

Output tensor after applying the multiplication operation.

Return type:

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

backward(grad: LazyBuffer) → LazyBuffer[source]

Backward method for negating a given gradient during the backward pass.

Parameters:: grad (LazyBuffer) – Gradient to be negated.
Returns:: The negation of the input gradient grad.
Return type:: LazyBuffer

forward(x: LazyBuffer) → LazyBuffer[source]

Forward method for negating a given value during the forward pass.

Parameters:: x (LazyBuffer) – Input buffer to be negated.
Returns:: The negation of the input buffer x.
Return type:: LazyBuffer

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

backward(grad_output: LazyBuffer) → LazyBuffer[source]

Returns the gradient output without any modification.

Parameters:: grad_output (LazyBuffer) – The gradient tensor.
Returns:: The same as the input grad_output tensor.
Return type:: LazyBuffer

forward(x: LazyBuffer) → LazyBuffer[source]

Returns a contiguous version of the input tensor.

Parameters:: x (LazyBuffer) – The input tensor.
Returns:: A contiguous version of the input tensor.
Return 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

forward(x: LazyBuffer) → LazyBuffer[source]

Apply the sine function to input LazyBuffer and return result.

Parameters:: x (LazyBuffer) – Input buffer for which the sine operation is applied.
Returns:: Output buffer after applying the sine 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]]

forward(x: LazyBuffer, y: LazyBuffer) → LazyBuffer[source]

Forward method to be implemented by subclasses. This method is called when the function is called in forward mode.

Raises:: NotImplementedError – If not overridden in a subclass.

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.

forward[source]

Applies the forward pass of the operation.

Type:: method

backward[source]

Applies the backward pass of the operation.

Type:: method

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