Return to Answer

1. Subscripts processing:

   >>> expr = 'ij,jk->ki'
   >>> qry_expr, res_expr = expr.split('->')
   >>> inputs_expr = qry_expr.split(',')
   >>> inputs_expr, res_expr
   (['ij', 'jk'], 'ki')
   

2. Find the unique keys (labels) in the input subscripts:

   >>> keys = set([(key, size) for keys, input in zip(inputs_expr, inputs) 
                  for key, size in list(zip(keys, input.shape))])
   {('i', 3), ('j', 4), ('k', 2)}
   

   We should be checking for constraints (as well as in the output subscript)! Using set is a bad idea but it will work for the purpose of this example.

3. Get the associated sizes (used to initialize the output array) and construct the ranges (used to create our domain of iteration):

   >>> sizes = dict(keys)
   {'i': 3, 'j': 4, 'k': 2}
   
   >>> ranges = [range(size) for _, size in keys]
   [range(0, 2), range(0, 3), range(0, 4)]
   

4. We need an list containing the keys (labels):

   >>> to_key = sizes.keys()
   ['k', 'i', 'j']
   

5. Compute the cartesian product of the ranges

   >>> domain = product(*ranges)
   

   Note: [itertools.product][1] returns an iterator which gets consumed over time.

6. Initialize the output tensor as:

   >>> res = np.zeros([sizes[key] for key in res_expr])
   

7. We will be looping over domain:

   >>> for indices in domain:
   ...     pass
   

   For each iteration, indices will contain the values on each axis. In our example, that would provide i, j, and k as a tuple: (k, i, j). For each input (A and B) we need to determine which component to fetch. That's A[i, j] and B[j, k], yes! However, we don't have variables i, j, and k, literally speaking.

   We can zip indices with to_key to create a mapping between each key (label) and its current value:

   >>> vals = dict(zip(to_key, indices))
   

   To get the coordinates for the output array, we use vals and loop over the keys: [vals[key] for key in res_expr]. However, to use these to index the output array, we need to wrap it with tuple and zip to separate the indices along each axis:

   >>> res_ind = tuple(zip([vals[key] for key in res_expr]))
   

   Same for the input indices (although there can be several):

   >>> inputs_ind = [tuple(zip([vals[key] for key in expr])) for expr in inputs_expr]
   

8. We will use a itertools.reduce to compute the product of all contributing components:

   >>> def reduce_mult(L):
   ...     return reduce(lambda x, y: x*y, L)
   

9. Overall the loop over the domain looks like:

   >>> for indices in domain:
   ...     vals = {k: v for v, k in zip(indices, to_key)}
   ...     res_ind = tuple(zip([vals[key] for key in res_expr]))
   ...     inputs_ind = [tuple(zip([vals[key] for key in expr])) 
   ...                     for expr in inputs_expr]
   ...
   ...     res[res_ind] += reduce_mult([M[i] for M, i in zip(inputs, inputs_ind)])
   

1. Subscripts processing:

   >>> expr = 'ij,jk->ki'
   >>> qry_expr, res_expr = expr.split('->')
   >>> inputs_expr = qry_expr.split(',')
   >>> inputs_expr, res_expr
   (['ij', 'jk'], 'ki')
   

2. Find the unique keys (labels) in the input subscripts:

   >>> keys = set([(key, size) for keys, input in zip(inputs_expr, inputs) 
                  for key, size in list(zip(keys, input.shape))])
   {('i', 3), ('j', 4), ('k', 2)}
   

   We should be checking for constraints (as well as in the output subscript)! Using set is a bad idea but it will work for the purpose of this example.

3. Get the associated sizes (used to initialize the output array) and construct the ranges (used to create our domain of iteration):

   >>> sizes = dict(keys)
   {'i': 3, 'j': 4, 'k': 2}
   
   >>> ranges = [range(size) for _, size in keys]
   [range(0, 2), range(0, 3), range(0, 4)]
   

4. We need an list containing the keys (labels):

   >>> to_key = sizes.keys()
   ['k', 'i', 'j']
   

5. Compute the cartesian product of the ranges

   >>> domain = product(*ranges)
   

   Note: [itertools.product][1] returns an iterator which gets consumed over time.

6. Initialize the output tensor as:

   >>> res = np.zeros([sizes[key] for key in res_expr])
   

7. We will be looping over domain:

   >>> for indices in domain:
   ...     pass
   

   For each iteration, indices will contain the values on each axis. In our example, that would provide i, j, and k as a tuple: (k, i, j). For each input (A and B) we need to determine which component to fetch. That's A[i, j] and B[j, k], yes! However, we don't have variables i, j, and k, literally speaking.

   We can zip indices with to_key to create a mapping between each key (label) and its current value:

   >>> vals = dict(zip(to_key, indices))
   

   To get the coordinates for the output array, we use vals and loop over the keys: [vals[key] for key in res_expr]. However, to use these to index the output array, we need to wrap it with tuple and zip to separate the indices along each axis:

   >>> res_ind = tuple(zip([vals[key] for key in res_expr]))
   

   Same for the input indices (although there can be several):

   >>> inputs_ind = [tuple(zip([vals[key] for key in expr])) for expr in inputs_expr]
   

8. We will use a functools.reduce to compute the product of all contributing components:

   >>> def reduce_mult(L):
   ...     return reduce(lambda x, y: x*y, L)
   

9. Overall the loop over the domain looks like:

   >>> for indices in domain:
   ...     vals = {k: v for v, k in zip(indices, to_key)}
   ...     res_ind = tuple(zip([vals[key] for key in res_expr]))
   ...     inputs_ind = [tuple(zip([vals[key] for key in expr])) 
   ...                     for expr in inputs_expr]
   ...
   ...     res[res_ind] += reduce_mult([M[i] for M, i in zip(inputs, inputs_ind)])
   

np.einsum is a general matrix operation tool. The string given contains the subscripts which are labels representing axes. I like to think of it as your operation definition. The subscripts provide two apparent constraints:

which components of the input arrays will contribute to component [k, i] of the output array.

Here is, somewhat, a minimal implementation of np.einsum in Python. For some, this might help to understand what's really going on under the hood. I find it as intuitive as the usage of np.einsum itself.

As we go along I will keep referring to the previous example. And define inputs as [A, B].

np.einsum can actually take more than two inputs! In the following, we will focus on the general case: n inputs and n input subscripts. The main goal is to find the domain of study, i.e. the cartesian product of all possible ranges.

We can't rely on manually writing for loops, simply because we don't know how many we will need! The main idea is this: we need to find all unique labels (I will use key and keys to refer to them), find the corresponding array shape, then create ranges for each one, and compute the product of the ranges using itertools.product to get the domain of study.

1. Subscripts processing:

   >>> expr = 'ij,jk->ki'
   >>> qry_expr, res_expr = expr.split('->')
   >>> inputs_expr = qry_expr.split(',')
   >>> inputs_expr, res_expr
   (['ij', 'jk'], 'ki')
   

2. Find the unique keys (labels) in the input subscripts:

   >>> keys = set([(key, size) for keys, input in zip(inputs_expr, inputs) 
                  for key, size in list(zip(keys, input.shape))])
   {('i', 3), ('j', 4), ('k', 2)}
   

   We should be checking for constraints (as well as in the output subscript)! Using set is a bad idea but it will work for the purpose of this example.

3. We need an list containing the keys (labels):

   >>> to_key = [key for key, _ in keys]
   ['k', 'i', 'j']
   

4. Get the associated sizes (used to initialize the output array) and construct the ranges (used to create our domain of iteration):

   >>> sizes = {key: size for key, size in keys}
   {'i': 3, 'j': 4, 'k': 2}
   
   >>> ranges = [range(size) for _, size in keys]
   [range(0, 2), range(0, 3), range(0, 4)]
   

5. Compute the cartesian product of the ranges

   >>> domain = product(*ranges)
   

   Beware: itertools.product returns an iterator which gets consumed over time.

6. Initialize the output tensor as:

   >>> res = np.zeros([sizes[key] for key in res_expr])
   

7. We will be looping over domain:

   >>> for indices in domain:
   ...     pass
   

   For each iteration, indices will contain the values on each axis. For our example, that would values i, j, and k as a tuple: (k, i, j). As if we were inside the three for loops. For each input (A and B) we need to determine which component to fetch. That's A[i, j] and B[j, k], yes! But that won't cut it since we don't have variables i, j, and k (literally speaking).

   We can zip indices with to_key to create a mapping between each key (label) and its current value:

   >>> vals= {k: v for v, k in zip(indices, to_key)}
   

   To get the coordinates for the output array, we use vals and loop over the keys: [vals[key] for key in res_expr]. However, to use these to index the output array, we need to wrap it with tuple and zip to separate the indices along each axis:

   >>> res_ind = tuple(zip([vals[key] for key in res_expr]))
   

   Same for the input indices (although there can be several):

   >>> inputs_ind = [tuple(zip([vals[key] for key in expr])) for expr in inputs_expr]
   

8. We will use a itertools.reduce to compute the product of all contributing components:

   >>> def reduce_mult(L):
   ...     return reduce(lambda x, y: x*y, L)
   

9. Overall the loop over the domain looks like:

   >>> for indices in domain:
   ...     vals = {k: v for v, k in zip(indices, to_key)}
   ...     res_ind = tuple(zip([vals[key] for key in res_expr]))
   ...     inputs_ind = [tuple(zip([vals[key] for key in expr])) 
   ...                     for expr in inputs_expr]
   ...
   ...     res[res_ind] += reduce_mult([M[i] for M, i in zip(inputs, inputs_ind)])
   

*Phew*, that's pretty close to np.einsum('ij,jk->ki', A, B)!

You can see einsum as a generalized matrix summation operator. The string given contains the subscripts which are labels representing axes. I like to think of it as your operation definition. The subscripts provide two apparent constraints:

which components from the input arrays will contribute to component [k, i] of the output array.

Here is a minimal implementation of np.einsum in Python. This should help understand what is really going on under the hood.

As we go along I will keep referring to the previous example. Defining inputs as [A, B].

np.einsum can actually take more than two inputs. In the following, we will focus on the general case: n inputs and n input subscripts. The main goal is to find the domain of iteration, i.e. the cartesian product of all our ranges.

We can't rely on manually writing for loops, simply because we don't know how many there will be. The main idea is this: we need to find all unique labels (I will use key and keys to refer to them), find the corresponding array shape, then create ranges for each one, and compute the product of the ranges using itertools.product to get the domain of study.

1. Subscripts processing:

   >>> expr = 'ij,jk->ki'
   >>> qry_expr, res_expr = expr.split('->')
   >>> inputs_expr = qry_expr.split(',')
   >>> inputs_expr, res_expr
   (['ij', 'jk'], 'ki')
   

2. Find the unique keys (labels) in the input subscripts:

   >>> keys = set([(key, size) for keys, input in zip(inputs_expr, inputs) 
                  for key, size in list(zip(keys, input.shape))])
   {('i', 3), ('j', 4), ('k', 2)}
   

   We should be checking for constraints (as well as in the output subscript)! Using set is a bad idea but it will work for the purpose of this example.

3. Get the associated sizes (used to initialize the output array) and construct the ranges (used to create our domain of iteration):

   >>> sizes = dict(keys)
   {'i': 3, 'j': 4, 'k': 2}
   
   >>> ranges = [range(size) for _, size in keys]
   [range(0, 2), range(0, 3), range(0, 4)]
   

4. We need an list containing the keys (labels):

   >>> to_key = sizes.keys()
   ['k', 'i', 'j']
   

5. Compute the cartesian product of the ranges

   >>> domain = product(*ranges)
   

   Note: [itertools.product][1] returns an iterator which gets consumed over time.

6. Initialize the output tensor as:

   >>> res = np.zeros([sizes[key] for key in res_expr])
   

7. We will be looping over domain:

   >>> for indices in domain:
   ...     pass
   

   For each iteration, indices will contain the values on each axis. In our example, that would provide i, j, and k as a tuple: (k, i, j). For each input (A and B) we need to determine which component to fetch. That's A[i, j] and B[j, k], yes! However, we don't have variables i, j, and k, literally speaking.

   We can zip indices with to_key to create a mapping between each key (label) and its current value:

   >>> vals = dict(zip(to_key, indices))
   

   To get the coordinates for the output array, we use vals and loop over the keys: [vals[key] for key in res_expr]. However, to use these to index the output array, we need to wrap it with tuple and zip to separate the indices along each axis:

   >>> res_ind = tuple(zip([vals[key] for key in res_expr]))
   

   Same for the input indices (although there can be several):

   >>> inputs_ind = [tuple(zip([vals[key] for key in expr])) for expr in inputs_expr]
   

8. We will use a itertools.reduce to compute the product of all contributing components:

   >>> def reduce_mult(L):
   ...     return reduce(lambda x, y: x*y, L)
   

9. Overall the loop over the domain looks like:

   >>> for indices in domain:
   ...     vals = {k: v for v, k in zip(indices, to_key)}
   ...     res_ind = tuple(zip([vals[key] for key in res_expr]))
   ...     inputs_ind = [tuple(zip([vals[key] for key in expr])) 
   ...                     for expr in inputs_expr]
   ...
   ...     res[res_ind] += reduce_mult([M[i] for M, i in zip(inputs, inputs_ind)])
   

That's pretty close to what np.einsum('ij,jk->ki', A, B) returns!

As you will see later down, there are other constraints. For instance

3. labels in the output subscript must not appear more than once.

When looking at ij,jk->ki, you can think of it as

We will stick with the following definitions of A and B:

The output Z will have a shape of (B.shape[1], A.shape[0]) and could naively be constructed in the following way. Remember np.einsum is about accumulating contributions in the output array. We will start with a blank array Z:

Each (A[i,j], B[j,k]) pair is seen contributing to each Z[k, i] component.

np.einsum can actually take more than two inputs... Initially, I had planned to post a small example using only one input but decided otherwise. In the following, we will focus on the general case: n inputs and n input subscripts. The main goal is to find the domain of study which the cartesian product of all axis' ranges.

The domain of study is the cartesian product: range(0, 3) x range(0, 4) x range(0, 2).

As you will see later down, there are other constraints. For instance:

3. labels in the output subscript must not appear more than once.

4. labels in the output subscript must appear in the input subscripts.

When looking at ij,jk->ki, you can think of it as:

We will stick with operation ij,jk->ki, and the following definitions of A and B:

The output, Z, will have a shape of (B.shape[1], A.shape[0]) and could naively be constructed in the following way. Starting with a blank array for Z:

np.einsum is about accumulating contributions in the output array. Each (A[i,j], B[j,k]) pair is seen contributing to each Z[k, i] component.

np.einsum can actually take more than two inputs! In the following, we will focus on the general case: n inputs and n input subscripts. The main goal is to find the domain of study, i.e. the cartesian product of all possible ranges.

The domain of study is the cartesian product: range(0, 2) x range(0, 3) x range(0, 4).