contract_expression(subscripts, *shapes, **kwargs)¶
Generate a reusable expression for a given contraction with specific shapes, which can, for example, be cached.
- subscripts (str) – Specifies the subscripts for summation.
- shapes (sequence of integer tuples) – Shapes of the arrays to optimize the contraction for.
- constants (sequence of int, optional) – The indices of any constant arguments in
shapes, in which case the actual array should be supplied at that position rather than just a shape. If these are specified, then constant parts of the contraction between calls will be reused. Additionally, if a GPU-enabled backend is used for example, then the constant tensors will be kept on the GPU, minimizing transfers.
- kwargs – Passed on to
expr – Callable with signature
expr(*arrays, out=None, backend='numpy')where the array’s shapes should match
- The out keyword argument should be supplied to the generated expression rather than this function.
- The backend keyword argument should also be supplied to the generated expression. If numpy arrays are supplied, if possible they will be converted to and back from the correct backend array type.
- The generated expression will work with any arrays which have the same rank (number of dimensions) as the original shapes, however, if the actual sizes are different, the expression may no longer be optimal.
- Constant operations will be computed upon the first call with a particular backend, then subsequently reused.
>>> expr = contract_expression("ab,bc->ac", (3, 4), (4, 5)) >>> a, b = np.random.rand(3, 4), np.random.rand(4, 5) >>> c = expr(a, b) >>> np.allclose(c, a @ b) True
aas a constant:
>>> expr = contract_expression("ab,bc->ac", a, (4, 5), constants=) >>> expr <ContractExpression('[ab],bc->ac', constants=)>
>>> c = expr(b) >>> np.allclose(c, a @ b) True