Tensor operations and primitives (with Tensorflow)

For fluently reading and writing custom TensorFlow code (e.g., custom loss functions, custom layers, …) a good understanding of basic tensor operations is necessary. This article gives an overview of operations, which are commonly found in the wild.

Terminology

A tensor is a multi-dimensional data structure (basically another name for a multi-dimensional array). Tensors (tf.Tensor) have a fixed shape, which defines the number of dimensions (aka axis) and the number of items within each dimension. The total number of dimensions is defined as the rank of a tensor. The total number of items within the data structure states the size of the tensor. Furthermore the tensor has a defined datatype (e.g., tf.int32, tf.float32) and all items must be of this datatype.

The following code snippet creates a tensor with the size of 9 and the rank of 2. The shape is (3, 3) and the datatype is implicitly set to tf.int32.

tensor = tf.constant([
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
])

print('shape: {}'.format(tensor.shape))
print('datatype: {}'.format(tensor.dtype))
print('rank: {}'.format(tf.rank(tensor)))
print('size: {}'.format(tf.size(tensor)))

Output:

shape: (3, 3)
datatype: <dtype: 'int32'>
rank: 2
size: 9

Based on the rank different mathematical names are defined for the tensor data structure. That is, a tensor with rank 0 is called a scalar, a tensor with rank 1 is called a vector, a tensor with rank 2 is called a matrix. Everything which has a rank higher than 2 is a tensor.

Creating tensors

tf.constant creates a immutable tensor from a tensor like data structure (e.g., a python list):

tensor = tf.constant([
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
])
print(tensor)

Output:

tf.Tensor(
[[1 2 3]
 [4 5 6]
 [7 8 9]], shape=(3, 3), dtype=int32)

tf.range creates a tensor similarly to a python range definition, with a start, end (optional) and delta (optional, default=1) value.

tensor = tf.range(1, 8, 2)
print(tensor)

Output:

tf.Tensor([1 3 5 7], shape=(4,), dtype=int32)

tf.fill creates a tensor based on a specified shape and sets a given value for each item. tf.ones and tf.zeros are similar to tf.fill, however instead of an arbitrary value 1 or 0 is used. Furthermore tf.ones_like and tf.zeros_like create a tensor with 1 or 0 based on the shape of an existing tensor.

tensor = tf.fill([3, 2], 5)
print(tensor)

Output:

tf.Tensor(
[[5 5]
 [5 5]
 [5 5]], shape=(3, 2), dtype=int32)

Indexing and slicing

Similar to numpy or python data structures certain indices can be accessed or parts of a tensor can be selected by slicing.

Important syntax for indexing and slicing:

negative indices count backward from the end of the tensor dimension
slicing can be used with :
... includes all dimension, which is a shorthand for using : repeatingly

tensor = tf.constant([
    [
        [1, 2, 3],
        [4, 5, 6]
    ],
    [
        [3, 2, 1],
        [6, 5, 4]
    ]
])
print(tensor[:, :, -1], tensor[..., -1])

Output:

tf.Tensor(
[[3 6]
 [1 4]], shape=(2, 2), dtype=int32)
tf.Tensor(
[[3 6]
 [1 4]], shape=(2, 2), dtype=int32)

Mathematical/Arithmetic operations

There are many mathematical operations implemented for tensors. Some are also overloaded with respective operators.

In the following code snippet several element-wise operations are demonstrated. Generally, for element-wise operations the tensors require the same shape. Broadcasting, which is explained in the next section allows for element-wise operations which don’t require the same shape.

a = tf.constant([1, 2, 3])
b = tf.constant([4, 5, 6])

print(a/b, tf.truediv(a, b))
print(a+b, tf.add(a, b))
print(a-b, tf.subtract(a, b))
print(a*b, tf.multiply(a, b))

Output:

tf.Tensor([0.25 0.4  0.5 ], shape=(3,), dtype=float64) tf.Tensor([0.25 0.4  0.5 ], shape=(3,), dtype=float64)
tf.Tensor([5 7 9], shape=(3,), dtype=int32) tf.Tensor([5 7 9], shape=(3,), dtype=int32)
tf.Tensor([-3 -3 -3], shape=(3,), dtype=int32) tf.Tensor([-3 -3 -3], shape=(3,), dtype=int32)
tf.Tensor([ 4 10 18], shape=(3,), dtype=int32) tf.Tensor([ 4 10 18], shape=(3,), dtype=int32)

Also the matrix multiplication is implemented:

a = tf.constant([
    [1, 2, 3],
    [4, 5, 6]
])
b = tf.constant([
    [4],
    [6],
    [8]
])
print(a @ b, tf.matmul(a, b))

Output:

tf.Tensor(
[[40]
 [94]], shape=(2, 1), dtype=int32)

tf.Tensor(
[[40]
 [94]], shape=(2, 1), dtype=int32)

Broadcasting

Broadcasting enables element-wise operations of tensors with different shapes. Thereby the tensor with the smaller shape for a certain dimension is “streched”, in order to execute the operation.

a = tf.constant([1, 2, 3, 4])
b = tf.constant([5])
print(a + b)
print(a * b)

Output:

tf.Tensor([6 7 8 9], shape=(4,), dtype=int32)
tf.Tensor([ 5 10 15 20], shape=(4,), dtype=int32)

In the example above the tensor b is broadcast to fit the shape of tensor a. The tensor would be broadcast into the following structure: [5, 5, 5, 5].

Reshaping tensors

The memory layout of a tf.Tensor is organized in row-major ordering (C-Style). When reshaping a tensor the actual data in the memory is not changed, only the index is reorganized. Therefore the input and output tensors of a reshaping operation must have the same size.

The tf.reshape operation is taking a tensor as input and additionally a shape definition:

tensor = tf.constant([
    [
        [1., 6., 4.],
        [2., 5., 1.]
    ],
    [
        [3., 8., 3.],
        [4., 9., 8.]
    ],
])

# reshape: (2, 2, 3) >>> (12, )

result1 = tf.reshape(tensor, [12])
print(result1)

# reshape: (2, 2, 3) >>> (4, 1, 1, 3)

result2 = tf.reshape(tensor, [4, 1, 1, 3])
print(result2)

Output:

tf.Tensor([1. 6. 4. 2. 5. 1. 3. 8. 3. 4. 9. 8.], shape=(12,), dtype=float32)
tf.Tensor(
[[[[1. 6. 4.]]]
 [[[2. 5. 1.]]]
 [[[3. 8. 3.]]]
 [[[4. 9. 8.]]]], shape=(4, 1, 1, 3), dtype=float32)

The tf.expand_dims and the tf.squeeze operations add or remove dimensions of length 1.

Tiling

Create an output tensor by tiling an input tensor according to a given specification. The tf.tile operation requires an input tensor and a multiples specification. The multiples specification is a one-dimensional tensor and the length of the dimension must be the same as the rank of the input tensor.

tensor = tf.constant([
    [
        [1, 2, 3],
        [4, 5, 6]
    ]
])
print(tf.tile(tensor, tf.constant([1, 3, 2])))

Output:

tf.Tensor(
[[[1 2 3 1 2 3]
  [4 5 6 4 5 6]
  [1 2 3 1 2 3]
  [4 5 6 4 5 6]
  [1 2 3 1 2 3]
  [4 5 6 4 5 6]]], shape=(1, 6, 6), dtype=int32)

The tf.tile operation replicates the given tensor of shape (1, 2, 3) according to the multiples specification [1, 3, 2]. That is, axis 0 is replicated once, axis 1 is replicated three times and axis 2 is replicated twice.

Repeat

The tf.repeat operation repeats tensor items on a specified axis according to a specified pattern. The length of the repeat pattern must be of the same length as the selected axis in the input tensor shape.

tensor = tf.constant([
    [1, 2, 3],
    [3, 4, 5]
])
print(tf.repeat(tensor, [2, 3], 0))
print(tf.repeat(tensor, [2, 4, 1], 1))

Output:

tf.Tensor(
[[1 2 3]
 [1 2 3]
 [3 4 5]
 [3 4 5]
 [3 4 5]], shape=(5, 3), dtype=int32)

tf.Tensor(
[[1 1 2 2 2 2 3]
 [3 3 4 4 4 4 5]], shape=(2, 7), dtype=int32)

Concat

The tf.concat operation “merges” tensors along a specified axis. The axis could also be specified with a negative number (similar to python when using a negative index for a list).

t1 = [[1, 2, 3], [4, 5, 6]]
t2 = [[7, 8, 9], [10, 11, 12]]

# concat along the 0-axis

print(tf.concat([t1, t2], 0))

# concat along the 1-axis

print(tf.concat([t1, t2], 1))

Output:

tf.Tensor(
[[ 1  2  3]
 [ 4  5  6]
 [ 7  8  9]
 [10 11 12]], shape=(4, 3), dtype=int32)

tf.Tensor(
[[ 1  2  3  7  8  9]
 [ 4  5  6 10 11 12]], shape=(2, 6), dtype=int32)

Transpose

In the case of a 2-D tensor the tf.transpose operation performs a classic matrix transposition.

tensor = tf.constant([
    [1, 2, 3],
    [4, 5, 6]
])
print(tf.transpose(tensor))

Output:

tf.Tensor(
[[1 4]
 [2 5]
 [3 6]], shape=(3, 2), dtype=int32)

For tensors of a rank greater than 2 a permutation specification can be defined. This specification states the “position” of every index in the resulting tensor.

In the example a tensor with the shape (2, 2, 3) is transposed according to a permutation specification [0, 2, 1]. The original order of the axes is [0, 1, 2], thus the transposition “swaps” the second with the third axis.

tensor = tf.constant([[[1,  2,  3],
                       [4,  5,  6]],
                      [[7,  8,  9],
                       [10, 11, 12]]])
print(tf.transpose(tensor, [0, 2, 1]))

Output:

tf.Tensor(
[[[ 1  4]
  [ 2  5]
  [ 3  6]]

 [[ 7 10]
  [ 8 11]
  [ 9 12]]], shape=(2, 3, 2), dtype=int32)