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
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()
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 topt.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.