pytaco.einsum

pytaco.einsum(expr, *operands, out_format=None, dtype=None)

Evaluates the Einstein summation convention on the input operands.

The einsum summation convention employed here is very similar to NumPy’s but has some differences explained below.

Taco’s einsum can express a wide variety of linear algebra expressions in a simple fashion. Einsum can be used in implicit mode where no output indices are specified. In this mode, it follows the usual einstein summation convention to compute an output. In explicit mode, the user can force summation over specified subscript variables.

Note that this einsum parser is a subset of what taco can express. The full evaluate() supports a much larger range of possible expressions.

See the notes section for more details.

Parameters

expr: str

Specifies the subscripts for summation as a comma separated list of subscript variables. An implicit (Classical Einstein summation) is calculation is performed unless there is an explicit indicator ‘->’ included along with subscript labels specifying the output.

operands: list of array_like, tensors, scipy csr and scipy csc matrices

This specifies the operands for the computation. Taco will copy any NumPy arrays that are not stored in row-major or column-major format.

out_format: format, optional

The storage format of the output tensor. Note that unlike the other tensor functions the format here must be an actual format instance and cannot be a mode format. If the out_format is left to none, all the modes default to dense.

dtype: datatype, optional: The datatype of the output tensor.

Returns

out: tensor: The result of the einstein summation computation.

Warning

This differs from NumPy’s einsum in two important ways. The first is that the same subscript cannot appear more than once in a given operand. The second is that for sparse tensors, some expressions may require the user to explicitly transpose the tensors before passing them into einsum.

See also

evaluate()

Notes

einsum provides a succint way to represent a large number of tensor algebra expressions. A list of some possible operations along with some examples is presented below:

Sum axes of tensor tensor_sum()
Transpose tensors or transpose to dense tensors tensor.transpose()
Matrix multiplication and dot products dot() and matmul()
Tensor contractions tensordot()

The expr string is a comma separated list of subscript labels where each label corresponds to a dimension in the tensor. In implicit mode, repeated subscripts are summed. This means that pt.einsum('ij,jk', t1, t2) is matrix multiplication since the j indices are implicitly summed over.

In explicit mode, we can specify the output directly by using the ‘->’ identifier along with a list of the output labels. This means that expressions such as pt.einsum('i->', t) take the sum of all the elements in the vector and return a scalar. The above is equivalent to pt.tensor_sum(t)

einsum also allows for broadcasting using ellipses. For instance, to do a matrix multiply with only the right most dimensions of an order tensor, we can write pt.einsum("...ik,...kj->...ij")

Examples

>>> import numpy as np
>>> import pytaco as pt
>>> a = np.arange(25).reshape(5, 5)
>>> t = pt.tensor([5, 5], pt.csr)
>>> for i in range(5): t.insert([i, i], a[i, i])
>>> vec = np.arange(5)

# We can sum over any of the axes of the sparse tensor as follows:
>>> pt.einsum("ij->j", t).to_array() # defaults to dense vector
array([ 0.,  6., 12., 18., 24.], dtype=float32)

>>> pt.tensor_sum(t, axis=1).to_array()
array([ 0.,  6., 12., 18., 24.], dtype=float32)

# We can perform matrix multiply in implicit mode
# Note that is Sparse x Dense matrix multiply due to the formats

>>> pt.einsum("ij,jk", t, a, out_format=pt.csr).to_array()
array([[  0.,   0.,   0.,   0.,   0.],
       [ 30.,  36.,  42.,  48.,  54.],
       [120., 132., 144., 156., 168.],
       [270., 288., 306., 324., 342.],
       [480., 504., 528., 552., 576.]], dtype=float32)

# Create a SpMV kernel (since t is csr)
>>> pt.einsum("ij,j->j", t, vec).to_array()
array([ 0.,  6., 24., 54., 96.], dtype=float32)

# Equivalently, we could write
>>> np.array(pt.inner(t, vec, pt.dense), copy=False)
array([ 0.,  6., 24., 54., 96.], dtype=float32)

# Sum tensor elements
>>> pt.einsum("ij->", t)[0]
60.0
>>> pt.tensor_sum(t)[0]
60.0

The NumPy docs contain more examples and an in depth explanation of this notation can be found here.