How to Create Vector and Matrix Operations in TensorFlow
Machine learning applications store a great deal of data in vectors (onedimensional tensors) and matrices (twodimensional tensors). To process this data, TensorFlow provides many functions that operate on vectors and matrices. The following table lists these functions and provides a description of each.
Vector and Matrix Operations
Function  Description 
tensordot(a, b, axes, name=None) 
Returns the sum of products for the elements in the given axes 
cross(a, b, name=None) 
Returns the elementwise cross product 
diag(diagonal, name=None) 
Returns a matrix with the given diagonal values, other values set to zero 
trace(x, name=None) 
Returns the sum of the diagonal elements 
transpose(x, perm=None,

Switches rows and columns 
eye(num_rows, num_columns=None,

Creates an identity matrix with the given shape and data type 
matmul(a, b,

Returns the product of the two input matrices 
norm(tensor,

Returns the norm of the given axis of the input tensor with the specified order 
matrix_solve(A, b,

Returns the tensor x, such that Ax = b, where A is a matrix, and b is a vector 
qr(input, full_matrices=None,

Returns the eigenvectors and eigenvalues of the given matrix or matrices 
svd(tensor,

Factors the matrix into a unitary matrix, a diagonal matrix, and the conjugate transpose of the unitary matrix 
einsum(equation, *inputs) 
Executes a custom mathematical operation 
Of these functions, the two most common are tensordot
and matmul
. tensordot
returns the dot product of one or more axes of two input tensors. That is, tensordot
multiplies the corresponding elements of both tensors’ dimensions and returns the sum of the products.
The axes
parameter tells tensordot
which dimensions to process. If you set this parameter to a scalar, N, the function will access the last N axes of the first tensor and the first N axes of the second tensor. If you set axes
equal to a list or tensor, the first row identifies axes of the first tensor, and the second row identifies axes of the second tensor.
I frequently call tensordot
to compute the dot product of two onedimensional tensors. The following code shows what this looks like:
t1 = tf.constant([4., 3., 2.]) t2 = tf.constant([3., 2., 1.]) dot = tf.tensordot(t1, t2, 1) # 4*3 + 3*2 + 2*1 = 20
matmul
performs traditional matrix multiplication. That is, it multiplies rows of the first tensor by columns of the second tensor and returns a matrix containing the sums. The following code shows how this can be used:
t1 = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) t2 = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) dot = tf.matmul(t1, t2) # [[ 22. 28.], [ 49. 64.]]
My favorite function in the table is einsum
, which makes it possible to create and execute custom mathematical operations. The first parameter is a string that identifies the operation using a special format called the Einstein summation convention. This convention has a number of characteristics, including the following:
 The operation is assumed to have one or two inputs. If you provide two inputs, you must separate them with a comma.
 Dimensions of input and output matrices are represented by subscripts (usually
i
,j
, andk
). Input subscripts must be separated from output subscripts with the → symbol.  If an input’s subscript is repeated and no output subscripts are given, the operation performs addition. Therefore,
einsum('i,i', t1, t2)
computes the dot product of tensorst1
andt2
.  If an input’s subscript is repeated and output subscripts are given, the operation performs multiplication. Therefore,
einsum('i,i>i', t1, t2)
computes the elementwise product of tensorst1
andt2
.
The following code calls einsum
to transpose a matrix and multiply two matrices together:
m1 = tf.constant([[1, 2], [3, 4]]) m2 = tf.constant([[5, 6], [7, 8]]) e1 = tf.einsum'ij<ji', m1) # [[1, 3], [2, 4]] e2 = tf.einsum('ij,jk<ik', m1, m2) # [[19, 22], [43, 50]]
For a more complete discussion of the Einstein summation convention, try Samuel Prime’s presentation.