pytaco.evaluate

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

Evaluates the index notation expression on the input operands.

An output tensor may be optionally specified. In this case, the tensor should be given the expected output shape, format and dtype since the out_format and dtype fields will be ignored if an output tensor is seen.

Parameters

expr: str

Specifies an index expression as a string. This must be of the form res(i1, i2, ...) = expr. See the examples for a more specific example. Each object represented by a name in the string and is indexed by a variable for each dimension.

operands: list of tensors or array_like

Specifies the input tensors OR the input and output tensor. If the length of the list is equal to N - 1 where N is the number of terms in the input expression then it is assumed that no output tensor was specified and taco infers the output shape and uses the out_format and dtype passed in for the output tensor. If the length of the operands is equal to N, then the first tensor is assumed to be the output tensor and the out_format and dtype fields are ignored.

out_format: format, optional

The storage format of the output tensor if one was not explicitly provided. If left to None and no output tensor was provided then all modes default to dense.

dtype: datatype

The datatype of the output tensor. If left to None and no output tensor was explicitly specified then taco uses its promotion rules and sets the output to the datatype with the highest type.

Returns

output: tensor

The tensor calculated based on the string expression passed in. Even if taco detects that an output is specified, it will still return a reference to that tensor.

Notes

This provides a convenient way to express tensor expressions. It is identical to the Index Expression syntax with a few exceptions. There is no need to use t[None] when making expressions with scalars and square brackets are replaced with parenthesis. For example, in python we can represent matrix multiply as A[i, j] = B[i, k] * C[k, j] while the corresponding tensor expression would be A(i, j) = B(i, k) * C(k, j). Further, reductions in pythonic index expression notation would be expressed as A[None] = B[i, j] to sum all the elements of a matrix while the corresponding string would be A = B(i, j).

The string parser currently only supports +, -, /, and *. Thus, expressions involving other functions such as exp, tan etc, would have to be written using the pythonic expressions.

An input tensor is recognised by the parser by a name followed by a comma separated list of index variables in parenthesis. Thus A(i,j,k) represents an order 3 tensor with the name A. The names used in the expression are irrelevant since taco will match the operands with the terms in the expression in the same order they appear (which is why when specifying an output, the output tensor must appear first followed by its input).

As with index expressions, index variables appearing that are on the right hand side of the expression but not in the result are always summed.

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)

# Note that no output is specified.
# We can sum over any of the axes of the sparse tensor as follows:
>>> pt.evaluate("T(j) = A(i, j)", t).to_array() # defaults to dense vector
array([ 0.,  6., 12., 18., 24.], dtype=float32)

# Specify an output
>>> result = pt.tensor([5], pt.dense)
>>> pt.evaluate("T(j) = A(i, j)", result, t).to_array()
array([ 0.,  6., 12., 18., 24.], dtype=float32)
>>> result.to_array()
array([ 0.,  6., 12., 18., 24.], dtype=float32)

# We can perform addition and broadcast along a given axis
>>> pt.evaluate("T(i, j) = A(i, j) + B(j)", t, vec, out_format=pt.csr).to_array()
array([[ 0.,  1.,  2.,  3.,  4.],
       [ 0.,  7.,  2.,  3.,  4.],
       [ 0.,  1., 14.,  3.,  4.],
       [ 0.,  1.,  2., 21.,  4.],
       [ 0.,  1.,  2.,  3., 28.]], dtype=float32)

# Create a SpMV kernel (since t is csr)
>>> pt.evaluate("A(j) = M(i, j) * V(j)", t, vec).to_array()
array([ 0.,  6., 24., 54., 96.], dtype=float32)

# Sum tensor elements, note that names used don't matter
>>> pt.evaluate("S = C(i, j)", t)[0]
60.0

Examples of reductions along with computations. Note indices that appear of the right hand side but not on the left hand side get summed over. This means we can implement matrix multiplication as shown below:

>>> from scipy.sparse import csc_matrix
>>> mat  = np.arange(9).reshape(3, 3)
>>> mat2 = csc_matrix(np.triu(np.arange(6).reshape(3, 2)))

# Compute mat @ mat2 due to ordering of operands.
>>> res = pt.evaluate("T(i, j) = A(i, k) * B(k, j)", mat, mat2, out_format=pt.csr)
>>> numpy_res = np.matmul(mat, mat2.toarray())
>>> all(res == numpy_res)
True