einops is a library for performing tensor operations that are more readable. It offers a simple API, that works with most tensor frameworks like PyTorch, JAX and Tensorflow.
rearrange
- reader-friendly smart element reordering for multidimensional tensors.
- This operation includes functionality of transpose (axes permutation), reshape (view), squeeze, unsqueeze, stack, concatenate and other operations.
repeat
- allows reordering elements and repeating them in arbitrary combinations.
- This operation includes functionality of repeat, tile, broadcast functions.
reduce
- provides combination of reordering and reduction using reader-friendly notation.
einsum
- Repeating letters in different inputs means those values will be multiplied, and those products will be in the output.
- For example,
M = einops.einsum(A, B, "i j, i j -> i j")just corresponds to the elementwise productM = A * B(because ).
- Omitting a letter means that the axis will be summed over.
- For example, if
xis a 2D array with shape(I, J), theneinops.einsum(x, "i j -> i")will be a 1D array of lengthIcontaining the row sums ofx(we’re summing along thej-index, i.e. across rows).
- We can return the un-summed axes in any order.
- For example,
einops.einsum(x, "i j k -> k j i")does the same thing aseinops.rearrange(x, "i j k -> k j i").
- free indices: are indices specified in the output. are associated with outer loop.
- summation indices: all other indices, those that appear in input but NOT in output.
- after initializing outer loops we initialize out and after inner loop we set out
A = np.array([0, 1, 2])
B = np.array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
np.einsum('i,ij->i', A, B) # array([ 0, 22, 76])
i is free indice
j is summuation indice
out[0] = A[0]* B[0][0] ... A[0]* B[0][3]
out[1] = A[1]* B[0][0] ... A[2]* B[0][3]
out[2] = A[2]* B[0][0] ... A[2]* B[0][3]
for i in range(A.shape[0])
sum = 0
for j in range(B.shape[1])
sum += A[i] * B[i][j]
out[i][j] = sum
References: