Tensor and Einsum

Task 1: Dot Product

To calculate the dot product of two vectors for variable length, one simply needs to loop over the single shared dimension.

Single loop over the shared dimension

    for i in range(a.shape[0]):
        result += a[i] * b[i]

Task 2: Matrix-Matrix Multiplication

Calculating a matrix-matrix multiplication simply adds another dimension to the calculation. Therefore, calculating the matrix-matrix multiplication using loops can be simply done by adding two more loops for the m and n dimensions.

Three loops to calculate the result of the matrix-matrix multiplication.

    for i in range(m):
        for j in range(n):
            for l in range(k):
                C[i, j] += A[i, l] * B[l, j]

Reusing the dot product calculation from Task 1 can be achieved by replacing the innermost loop by the function call to the dot_product function.

Two loops iterating over the dot_product function call.

    for i in range(m):
        for j in range(n):
            C[i, j] += dot_product(A[i, :], B[:, j])

Task 3: Einsum

In order to properly calculate the einsum expression, we need one loop for each dimension. The loops have been ordered according to their indices. The further an index is on the right, the further it is nested.

Several nested loops to calculate the einsum expression.

    for a in range(size_a):
        for b in range(size_b):
            for c in range(size_c):
                for s in range(size_s):
                    for x in range(size_x):
                        for p in range(size_p):
                            for y in range(size_y):
                                C[a, b, c, x, y] += A[a, c, s, x, p] * B[b, s, p, y]

The nesting depth of the first calculation can be reduced by invoking the matmul_dot function from Task 2. We selected the innermost dimensions of the matrices A and B to pass to the matmul_dot function.

Calculating the einsum expression by calling the matmul_dot function.

    for a in range(size_a):
        for b in range(size_b):
            for c in range(size_c):
                for s in range(size_s):
                    C[a, b, c, :, :] += matmul_dot(A[a, c, s, :, :], B[b, s, :, :])

Task 4: Tensor Permutation and Reshaping

The experiments for the tensor permutation and reshaping are located in the permute_reshape.py file.

Assume we initially have the following tensor:

# Initial tensor
tensor([[ 1,  2,  3,  4],
        [ 5,  6,  7,  8],
        [ 9, 10, 11, 12],
        [13, 14, 15, 16]])

# Shape
torch.Size([4, 4])

# Stride
(4, 1)

# Contiguity
True

Permutation

torch.permute takes two input parameters:

the input tensor and
the permuted dimensions.

This function returns a view of the original tensor with the permuted dimensions. Alternatively, one can also call <tensor>.permute(<dims>).

After applying the torch.permute(0, 1) function, the initial tensor changes to:

# Tensor after permute
tensor([[ 1,  5,  9, 13],
        [ 2,  6, 10, 14],
        [ 3,  7, 11, 15],
        [ 4,  8, 12, 16]])

# Shape
torch.Size([4, 4])

# Stride
(1, 4)

# Contiguity
False

This shows that the memory layout does change after applying the torch.permute function. Furthermore, the resulting tensor is no longer stored continuously in memory.

Reshape

torch.reshape takes two input parameters:

the input tensor and
the new shape.

This function returns a tensor with the same data and number of elements as the original input tensor, simply with a new shape. If possible, the returned tensor comes in the shape of a view, otherwise it will be a copy. For example, contiguous inputs and inputs with compatible strides are reshaped without copying.

In comparison, if we apply the torch.reshape(1, -1) function, the tensor looks like:

# Tensor after reshape
tensor([[ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16]])

# Shape
torch.Size([1, 16])

# Stride
(16, 1)

# Contiguity
True

This shows that the order of the tensor elements is the same as the initial tensor. Therefore, the memory layout does not change after this torch.reshape function call.

View

Depending on the contiguity of the given tensor, the result of the torch.view operation differs from torch.reshape. We experimented with two different parameter assignments. If we apply torch.reshape(1, -1) and torch.view(1, -1) to the same initial tensor, everything works the same.

However, if we first transpose the tensor using torch.transpose(0, 1), we can still apply the torch.reshape(1, -1) operation, but we can’t apply the torch.view(1, -1) operation.

# Tensor after transpose
tensor([[ 1,  5,  9, 13],
        [ 2,  6, 10, 14],
        [ 3,  7, 11, 15],
        [ 4,  8, 12, 16]])

# Contiguity
Tensor contiguity: False

# Tensor after transpose and reshape
tensor([[ 1,  5,  9, 13,  2,  6, 10, 14,  3,  7, 11, 15,  4,  8, 12, 16]])

# Shape
torch.Size([1, 16])

# Stride
(16, 1)

# Contiguity
True

It is noticeable that the torch.reshape operation can be applied to non-contiguous tensors. Furthermore, after applying torch.reshape, the memory layout is contiguous again.

The reason why a permuted tensor can be reshaped is that the torch.reshape operation is not influenced by the memory layout of the given tensor. As the description in torch already mentions, torch.reshape can handle tensors differently. In the case of a permuted tensor, the reshape operation creates a copy of the permuted tensor. In the case of a freshly created tensor, the reshape operation does not create a copy but rather creates a view of the freshly created tensor.