Operators

Differentiable operations that form the building blocks of the computation graph. Each op implements a compute method (forward pass on NDArrays) and a gradient method (backward pass for autograd).

These are used internally by Tensor methods like +, @, .reshape(), etc.

from tiny_pytorch import Tensor, ops

x = Tensor([[1.0, 2.0], [3.0, 4.0]])
y = ops.Summation(axes=(1,))(x)  # [3.0, 7.0]

Tensor operations module for tiny-pytorch implementation.

This module provides a comprehensive collection of fundamental tensor operations that form the building blocks of the computational graph in tiny-pytorch. Each operation is implemented as a class that inherits from the TensorOp base class, with corresponding helper functions for easier usage.

The module includes element-wise operations, matrix operations, reduction operations, activation functions, and various mathematical functions commonly used in deep learning and neural network computations.

Key Features

Automatic differentiation support through gradient methods
Element-wise and scalar operations
Matrix operations (multiplication, transpose)
Reduction operations (summation, log-sum-exp)
Activation functions (ReLU, tanh)
Shape manipulation (reshape, broadcast, stack, split)
Convolutional operations
Memory-efficient operations with strided arrays

Classes:

TensorOp –

Base class for all tensor operations.
TensorTupleOp –

Base class for operations that return tensor tuples.
ScalarAdd –

Addition of a scalar to a tensor.
EWiseAdd –

Element-wise addition of two tensors.
ScalarMul –

Multiplication of a tensor by a scalar.
EWiseMul –

Element-wise multiplication of two tensors.
Negate –

Negation of a tensor.
ScalarPower –

Raising tensor elements to a scalar power.
EWisePower –

Element-wise power operation between two tensors.
ScalarDivide –

Division of a tensor by a scalar.
EWiseDivide –

Element-wise division of two tensors.
Reshape –

Reshaping a tensor to a new shape.
Summation –

Summing tensor elements along specified axes.
BroadcastTo –

Broadcasting a tensor to a larger shape.
Transpose –

Transposing a tensor along specified axes.
MatMul –

Matrix multiplication between two tensors.
Log –

Natural logarithm of tensor elements.
Exp –

Exponential of tensor elements.
ReLU –

Rectified Linear Unit activation function.
LogSumExp –

Log-sum-exp operation, commonly used in softmax computation.
Tanh –

Hyperbolic tangent activation function.
Stack –

Stack a sequence of arrays along a new axis.
Split –

Split a tensor along a specified axis.
Flip –

Reverse the order of elements along specified axes.
Dilate –

Insert zeros between elements along specified axes.
UnDilate –

Remove zeros inserted by dilation along specified axes.
Conv –

2D convolution operation.

Functions:

add_scalar –

Add a scalar to a tensor.
add –

Add two tensors element-wise.
mul_scalar –

Multiply a tensor by a scalar.
multiply –

Multiply two tensors element-wise.
negate –

Negate a tensor.
power_scalar –

Raise tensor elements to a scalar power.
power –

Element-wise power operation.
divide_scalar –

Divide a tensor by a scalar.
divide –

Element-wise division of tensors.
reshape –

Reshape a tensor.
summation –

Sum tensor elements along specified axes.
broadcast_to –

Broadcast tensor to a larger shape.
transpose –

Transpose tensor axes.
matmul –

Matrix multiplication.
log –

Natural logarithm.
exp –

Exponential function.
relu –

ReLU activation function.
logsumexp –

Log-sum-exp operation.
tanh –

Hyperbolic tangent function.
stack –

Stack a sequence of arrays along a new axis.
split –

Split a tensor along a specified axis.
flip –

Reverse the order of elements along specified axes.
dilate –

Insert zeros between elements along specified axes.
undilate –

Remove zeros inserted by dilation along specified axes.
conv –

2D convolution operation.

Notes

All operations support automatic differentiation through their gradient methods, making them suitable for building and training neural networks. The operations are designed to work efficiently with the NDArray backend system and support multiple devices (CPU, CUDA, NumPy).

Examples:

>>> import tiny_pytorch as tp
>>> x = tp.Tensor([1, 2, 3])
>>> y = tp.Tensor([4, 5, 6])
>>> z = tp.ops.add(x, y)  # Element-wise addition
>>> w = tp.ops.matmul(x, y)  # Matrix multiplication

`BroadcastTo`

Bases: TensorOp

Broadcast a tensor to a larger shape.

Parameters:

shape (tuple) –

Target shape to broadcast to.

Methods:

compute –

Compute the broadcast operation.
gradient –

Compute the gradient of the operation.

`Conv`

Bases: TensorOp

2D convolution operation between input tensor and kernel.

This operation performs 2D convolution between an input tensor and a kernel tensor. The input is expected to be in NHWC format (batch, height, width, channels) and the kernel in KKCC format (kernel_height, kernel_width, input_channels, output_channels).

Parameters:

stride (int, default: 1 ) –

The stride of the convolution. Default is 1.
padding (int, default: 0 ) –

The amount of padding to apply to the input. Default is 0.

Methods:

compute –

Compute the 2D convolution operation using im2col and matrix multiplication.
gradient –

Compute the gradient with respect to both input and kernel tensors.

Notes

Input tensor A should have shape (N, H, W, C_in)
Kernel tensor B should have shape (K, K, C_in, C_out) where K is the kernel size
Output tensor will have shape (N, out_H, out_W, C_out)
Uses im2col transformation for efficient computation

`Dilate`

Bases: TensorOp

Dilate a tensor by inserting zeros between elements along specified axes.

This operation inserts zeros between elements along the specified axes, effectively increasing the size of the tensor in those dimensions. This is commonly used in convolutional neural networks for dilated convolutions.

Parameters:

axes (tuple[int, ...]) –

The axes along which to apply dilation. Each axis index must be valid for the tensor's dimensions.
dilation (int) –

The dilation factor. For each element in the original tensor, dilation zeros will be inserted after it along the specified axes.

Methods:

compute –

Compute the dilation operation on the input NDArray.
gradient –

Compute the gradient of the dilation operation (returns undilated gradient).

Examples:

>>> x = Tensor([[1, 2], [3, 4]])
>>> Dilate((0,), 1)(x)
Tensor([[1, 2], [0, 0], [3, 4]])
>>> Dilate((1,), 1)(x)
Tensor([[1, 0, 2], [3, 0, 4]])
>>> Dilate((0, 1), 1)(x)
Tensor([[1, 0, 2], [0, 0, 0], [3, 0, 4]])

`EWiseAdd`

Bases: TensorOp

Element-wise addition of two tensors.

Methods:

compute –

Compute element-wise addition.
gradient –

Compute the gradient with respect to both inputs.

`EWiseDivide`

Bases: TensorOp

Element-wise division of two tensors.

Methods:

compute –

Compute element-wise division.
gradient –

Compute the gradient with respect to both inputs.

`EWiseMul`

Bases: TensorOp

Element-wise multiplication of two tensors.

Methods:

compute –

Compute element-wise multiplication.
gradient –

Compute the gradient with respect to both inputs.

`EWisePower`

Bases: TensorOp

Element-wise power operation between two tensors.

Methods:

compute –

Compute element-wise power operation.
gradient –

Compute the gradient with respect to both inputs.

`Exp`

Bases: TensorOp

Exponential of tensor elements.

Methods:

compute –

Compute exponential.
gradient –

Compute the gradient of the operation.

`Flip`

Bases: TensorOp

Reverse (flip) the order of elements in a tensor along the specified axes.

Parameters:

axes (tuple[int, ...] or None, default: None ) –

Axes along which to flip the tensor. Each axis index must be valid for the tensor's dimensions. If None, flip over all axes (reverse the tensor in every dimension).

Methods:

compute –

Compute the flip operation on the input NDArray.
gradient –

Compute the gradient of the flip operation (flip the gradient along the same axes).

Raises:

AxisError –

If the number of axes is greater than the number of dimensions, or if any axis is out of bounds.

Examples:

>>> x = Tensor([[1, 2], [3, 4]])
>>> Flip((0,))(x)
Tensor([[3, 4], [1, 2]])
>>> Flip((1,))(x)
Tensor([[2, 1], [4, 3]])
>>> Flip((0, 1))(x)
Tensor([[4, 3], [2, 1]])
>>> Flip()(x)
Tensor([[4, 3], [2, 1]])

`Log`

Bases: TensorOp

Natural logarithm of tensor elements.

Methods:

compute –

Compute natural logarithm.
gradient –

Compute the gradient of the operation.

`LogSumExp`

Bases: TensorOp

Log-sum-exp operation, commonly used in softmax computation.

Parameters:

axes (tuple or None, default: None ) –

Axes along which to perform the operation. If None, use all axes.

Methods:

compute –

Compute log-sum-exp operation.
gradient –

Compute the gradient of the operation.

`MakeTensorTuple`

Bases: TensorOp

Pack multiple tensors into a TensorTuple.

`compute(*args)`

Pack input arrays into a Python tuple.

Parameters:

*args (NDArray, default: () ) –

Arrays to pack.

Returns:

tuple[NDArray] –

Tuple of input arrays.

`gradient(out_grad, node)`

Return each element of out_grad as a separate gradient.

Parameters:

out_grad (TensorTuple) –

Gradient of the output tuple.
node (Tensor) –

Output node of this operation.

Returns:

tuple[Tensor] –

Gradient with respect to each input tensor.

`MatMul`

Bases: TensorOp

Matrix multiplication between two tensors.

Methods:

compute –

Compute matrix multiplication.
gradient –

Compute the gradient with respect to both inputs.

`Negate`

Bases: TensorOp

Negate a tensor element-wise.

Methods:

compute –

Compute the negation operation.
gradient –

Compute the gradient of the operation.

`ReLU`

Bases: TensorOp

Rectified Linear Unit activation function.

Methods:

compute –

Compute ReLU activation.
gradient –

Compute the gradient of the operation.

`Reshape`

Bases: TensorOp

Reshape a tensor to a new shape.

Parameters:

shape (tuple) –

The target shape for the tensor.

Methods:

compute –

Compute the reshape operation.
gradient –

Compute the gradient of the operation.

`ScalarAdd`

Bases: TensorOp

Add a scalar to a tensor.

Parameters:

scalar (float) –

The scalar value to add to the tensor.

Methods:

compute –

Compute the scalar addition operation.
gradient –

Compute the gradient of the operation.

`ScalarDivide`

Bases: TensorOp

Divide a tensor by a scalar.

Parameters:

scalar (float) –

The scalar value to divide by.

Methods:

compute –

Compute the scalar division.
gradient –

Compute the gradient of the operation.

`ScalarMul`

Bases: TensorOp

Multiply a tensor by a scalar.

Parameters:

scalar (float) –

The scalar value to multiply with the tensor.

Methods:

compute –

Compute the scalar multiplication.
gradient –

Compute the gradient of the operation.

`ScalarPower`

Bases: TensorOp

Raise tensor elements to a scalar power.

Parameters:

scalar (float) –

The power to raise tensor elements to.

Methods:

compute –

Compute the power operation.
gradient –

Compute the gradient of the operation.

`Split`

Bases: TensorOp

Split a tensor along an axis into a tuple of tensors.

This operation is the inverse of Stack. It splits a tensor along a specified axis into multiple tensors, each with one less dimension than the input tensor.

Parameters:

axis (int) –

The axis along which to split the tensor. The axis dimension will be removed from each resulting tensor.

Methods:

compute –

Split the input array along the specified axis.
gradient –

Compute the gradient of the split operation (returns stack of out_grad tensors).

`Stack`

Bases: TensorOp

Stack a sequence of arrays along a new axis.

Parameters:

axis (int) –

The axis along which to stack. The new axis will be inserted at this position in the result array shape.

Methods:

compute –

Stack the input arrays along the specified axis.
gradient –

Compute the gradient of the stack operation (returns split of out_grad along axis).

`Summation`

Bases: TensorOp

Sum tensor elements along specified axes.

Parameters:

axes (tuple or None, default: None ) –

Axes along which to perform summation. If None, sum over all axes.

Methods:

compute –

Compute the summation operation.
gradient –

Compute the gradient of the operation.

`Tanh`

Bases: TensorOp

Hyperbolic tangent activation function.

Methods:

compute –

Compute hyperbolic tangent.
gradient –

Compute the gradient of the operation.

`Transpose`

Bases: TensorOp

Transpose a tensor along specified axes.

Parameters:

axes (tuple or None, default: None ) –

Permutation of the dimensions. If None, reverse the last two dimensions.

Methods:

compute –

Compute the transpose operation.
gradient –

Compute the gradient of the operation.

`TupleGetItem`

Bases: TensorOp

Extract a single tensor from a TensorTuple by index.

`compute(a)`

Extract element at self.index from the input tuple.

Parameters:

a (tuple[NDArray]) –

Input tuple of arrays.

Returns:

NDArray –

The array at the specified index.

`gradient(out_grad, node)`

Return gradient for the indexed element; zeros for all others.

Parameters:

out_grad (Tensor) –

Gradient of the output.
node (Tensor) –

Output node of this operation.

Returns:

TensorTuple –

Tuple with out_grad at the indexed position and zeros elsewhere.

`UnDilate`

Bases: TensorOp

Undilate a tensor by removing zeros inserted by dilation along specified axes.

This operation is the inverse of Dilate. It removes the zeros that were inserted during dilation, effectively reducing the size of the tensor in those dimensions. This is commonly used in convolutional neural networks for dilated convolutions.

Parameters:

axes (tuple[int, ...]) –

The axes along which to apply undilation. Each axis index must be valid for the tensor's dimensions.
dilation (int) –

The dilation factor that was used in the original Dilate operation.

Methods:

compute –

Compute the undilation operation on the input NDArray.
gradient –

Compute the gradient of the undilation operation (returns dilated gradient).

Examples:

>>> x = Tensor([[1, 0, 2], [0, 0, 0], [3, 0, 4]])
>>> UnDilate((0,), 1)(x)
Tensor([[1, 2], [3, 4]])
>>> UnDilate((1,), 1)(x)
Tensor([[1, 2], [3, 4]])
>>> UnDilate((0, 1), 1)(x)
Tensor([[1, 2], [3, 4]])

`add(a, b)`

Add two tensors element-wise.

Parameters:

a (Tensor) –

First input tensor.
b (Tensor) –

Second input tensor.

Returns:

Tensor –

Element-wise sum of the input tensors.

`add_scalar(a, scalar)`

Add a scalar value to a tensor.

Parameters:

a (Tensor) –

Input tensor.
scalar (float) –

Scalar value to add.

Returns:

Tensor –

A new tensor with the scalar added to each element.

`broadcast_to(a, shape)`

Broadcast a tensor to a larger shape.

Parameters:

a (Tensor) –

Input tensor.
shape (tuple) –

Target shape to broadcast to.

Returns:

Tensor –

Broadcasted tensor with the specified shape.

`conv(a, b, stride=1, padding=1)`

Perform 2D convolution between input tensor and kernel.

This function performs 2D convolution between an input tensor and a kernel tensor. The input is expected to be in NHWC format (batch, height, width, channels) and the kernel in KKCC format (kernel_height, kernel_width, input_channels, output_channels).

Parameters:

a (Tensor) –

Input tensor with shape (N, H, W, C_in) in NHWC format.
b (Tensor) –

Kernel tensor with shape (K, K, C_in, C_out) where K is the kernel size.
stride (int, default: 1 ) –

The stride of the convolution. Default is 1.
padding (int, default: 1 ) –

The amount of padding to apply to the input. Default is 1.

Returns:

Tensor –

Convolved tensor with shape (N, out_H, out_W, C_out) where: - out_H = (H + 2padding - K) // stride + 1 - out_W = (W + 2padding - K) // stride + 1

Notes

Uses im2col transformation for efficient computation
Supports automatic differentiation through gradient computation
Kernel must be square (K x K)
Input and kernel channel dimensions must match

Examples:

>>> from tiny_pytorch import init
>>> x = init.randn(1, 32, 32, 3)  # 1 batch, 32x32 image, 3 channels
>>> kernel = init.randn(3, 3, 3, 16)  # 3x3 kernel, 3 input channels, 16 output channels
>>> result = conv(x, kernel, stride=1, padding=1)
>>> result.shape  # (1, 32, 32, 16)
(1, 32, 32, 16)

`dilate(a, axes, dilation)`

Dilate a tensor by inserting zeros between elements along specified axes.

This function inserts zeros between elements along the specified axes, effectively increasing the size of the tensor in those dimensions. This is commonly used in convolutional neural networks for dilated convolutions.

Parameters:

a (Tensor) –

Input tensor to be dilated.
axes (tuple[int, ...]) –

The axes along which to apply dilation. Each axis index must be valid for the tensor's dimensions.
dilation (int) –

The dilation factor. For each element in the original tensor, dilation zeros will be inserted after it along the specified axes.

Returns:

Tensor –

A dilated tensor with zeros inserted along the specified axes.

Examples:

>>> x = Tensor([[1, 2], [3, 4]])
>>> dilate(x, (0,), 1)
Tensor([[1, 2], [0, 0], [3, 4]])
>>> dilate(x, (1,), 1)
Tensor([[1, 0, 2], [3, 0, 4]])
>>> dilate(x, (0, 1), 1)
Tensor([[1, 0, 2], [0, 0, 0], [3, 0, 4]])

`divide(a, b)`

Divide two tensors element-wise.

Parameters:

a (Tensor) –

Numerator tensor.
b (Tensor) –

Denominator tensor.

Returns:

Tensor –

Element-wise division of the input tensors.

`divide_scalar(a, scalar)`

Divide a tensor by a scalar value.

Parameters:

a (Tensor) –

Input tensor.
scalar (float) –

Scalar value to divide by.

Returns:

Tensor –

A new tensor with each element divided by the scalar.

`exp(a)`

Compute the exponential of tensor elements.

Parameters:

a (Tensor) –

Input tensor.

Returns:

Tensor –

Exponential of input tensor elements.

`flip(a, axes=None)`

Reverse (flip) the order of elements in a tensor along the specified axes.

Parameters:

a (Tensor) –

Input tensor to be flipped.
axes (tuple[int, ...] or None, default: None ) –

Axes along which to flip the tensor. Each axis index must be valid for the tensor's dimensions. If None, flip over all axes (reverse the tensor in every dimension).

Returns:

Tensor –

A tensor with the entries reversed along the specified axes.

Raises:

AxisError –

If the number of axes is greater than the number of dimensions, or if any axis is out of bounds.

Examples:

>>> x = Tensor([[1, 2], [3, 4]])
>>> flip(x, (0,))
Tensor([[3, 4], [1, 2]])
>>> flip(x, (1,))
Tensor([[2, 1], [4, 3]])
>>> flip(x, (0, 1))
Tensor([[4, 3], [2, 1]])
>>> flip(x)
Tensor([[4, 3], [2, 1]])

`log(a)`

Compute the natural logarithm of tensor elements.

Parameters:

a (Tensor) –

Input tensor.

Returns:

Tensor –

Natural logarithm of input tensor elements.

`logsumexp(a, axes=None)`

Compute log-sum-exp along specified axes.

Parameters:

a (Tensor) –

Input tensor.
axes (tuple or None, default: None ) –

Axes along which to perform the operation. If None, use all axes.

Returns:

Tensor –

Result of log-sum-exp operation.

`make_tuple(*args)`

Pack tensors into a TensorTuple.

Parameters:

*args (Tensor, default: () ) –

Tensors to pack.

Returns:

TensorTuple –

Tuple of input tensors.

`matmul(a, b)`

Perform matrix multiplication between two tensors.

Parameters:

a (Tensor) –

First input tensor.
b (Tensor) –

Second input tensor.

Returns:

Tensor –

Result of matrix multiplication.

`mul_scalar(a, scalar)`

Multiply a tensor by a scalar value.

Parameters:

a (Tensor) –

Input tensor.
scalar (float) –

Scalar value to multiply with.

Returns:

Tensor –

A new tensor with each element multiplied by the scalar.

`multiply(a, b)`

Multiply two tensors element-wise.

Parameters:

a (Tensor) –

First input tensor.
b (Tensor) –

Second input tensor.

Returns:

Tensor –

Element-wise product of the input tensors.

`negate(a)`

Negate a tensor element-wise.

Parameters:

a (Tensor) –

Input tensor.

Returns:

Tensor –

A new tensor with each element negated.

`power(a, b)`

Raise elements of one tensor to powers specified by another tensor.

Parameters:

a (Tensor) –

Base tensor.
b (Tensor) –

Exponent tensor.

Returns:

Tensor –

Element-wise power operation result.

`power_scalar(a, scalar)`

Raise tensor elements to a scalar power.

Parameters:

a (Tensor) –

Input tensor.
scalar (float) –

Power to raise elements to.

Returns:

Tensor –

A new tensor with each element raised to the given power.

`relu(a)`

Apply Rectified Linear Unit (ReLU) activation function.

Parameters:

a (Tensor) –

Input tensor.

Returns:

Tensor –

Tensor with ReLU activation applied.

`reshape(a, shape)`

Reshape a tensor to a new shape.

Parameters:

a (Tensor) –

Input tensor.
shape (tuple) –

Target shape for the tensor.

Returns:

Tensor –

A new tensor with the specified shape.

`split(a, axis)`

Split a tensor along an axis into a tuple of tensors.

This function splits a tensor along a specified axis into multiple tensors. Each resulting tensor has one less dimension than the input tensor.

Parameters:

a (Tensor) –

Input tensor to split.
axis (int) –

The axis along which to split the tensor. The axis dimension will be removed from each resulting tensor.

Returns:

TensorTuple –

A tuple of tensors, each with the specified axis dimension removed. The number of tensors in the tuple equals the size of the input tensor along the specified axis.

Examples:

>>> x = Tensor([[1, 2, 3], [4, 5, 6]])
>>> result = split(x, axis=0)
>>> len(result)  # Returns 2 tensors
2
>>> result[0].shape  # Each tensor has shape (3,)
(3,)

`stack(arrays, axis)`

Stack a sequence of tensors along a new axis.

Parameters:

arrays (list of Tensor) –

Sequence of tensors to stack. All tensors must have the same shape.
axis (int) –

The axis along which to stack. The new axis will be inserted at this position in the result tensor shape.

Returns:

Tensor –

The stacked tensor with one more dimension than the input tensors.

`summation(a, axes=None)`

Sum tensor elements along specified axes.

Parameters:

a (Tensor) –

Input tensor.
axes (tuple or None, default: None ) –

Axes along which to perform summation. If None, sum over all axes.

Returns:

Tensor –

Sum of elements along specified axes.

`tanh(a)`

Compute the hyperbolic tangent of tensor elements.

Parameters:

a (Tensor) –

Input tensor.

Returns:

Tensor –

Hyperbolic tangent of input tensor elements.

`transpose(a, axes=None)`

Transpose a tensor along specified axes.

Parameters:

a (Tensor) –

Input tensor.
axes (tuple or None, default: None ) –

Permutation of the dimensions. If None, reverse the dimensions.

Returns:

Tensor –

Transposed tensor.

`tuple_get_item(value, index)`

Extract a single tensor from a TensorTuple by index.

Parameters:

value (TensorTuple) –

Input tuple of tensors.
index (int) –

Index of the tensor to extract.

Returns:

Tensor –

The tensor at the specified index.

`undilate(a, axes, dilation)`

Undilate a tensor by removing zeros inserted by dilation along specified axes.

This function is the inverse of dilate. It removes the zeros that were inserted during dilation, effectively reducing the size of the tensor in those dimensions. This is commonly used in convolutional neural networks for dilated convolutions.

Parameters:

a (Tensor) –

Input tensor to be undilated.
axes (tuple[int, ...]) –

The axes along which to apply undilation. Each axis index must be valid for the tensor's dimensions.
dilation (int) –

The dilation factor that was used in the original dilate operation.

Returns:

Tensor –

An undilated tensor with zeros removed along the specified axes.

Examples:

>>> x = Tensor([[1, 0, 2], [0, 0, 0], [3, 0, 4]])
>>> undilate(x, (0,), 1)
Tensor([[1, 2], [3, 4]])
>>> undilate(x, (1,), 1)
Tensor([[1, 2], [3, 4]])
>>> undilate(x, (0, 1), 1)
Tensor([[1, 2], [3, 4]])