Interface
The @tensor macro rewrites tensor operations in terms of basic building blocks, such that any tensor type that implements the following interface can be supported. In these methods, C will indicate an output tensor which is changed in-place, while A and B denote input tensors that are left unaltered. pC, pA and pB denote an Index2Tuple, a tuple of two tuples that represents a permutation and partition of the original tensor indices. Finally, conjA and conjB are symbols that are used to indicate if the input tensor should be conjugated (:C) or used as-is (:N).
Operations
The three primitive tensor operations have already been described in the previous section, and correspond to
All other functions described in the previous section are implemented in terms of those. Hence, those are the only methods that need to be overloaded to support e.g. a new tensor type of to implement a new backend. There is one more necessary tensor operation, which is to convert back from a rank zero tensor to a scalar quantity.
TensorOperations.tensorscalar — Functiontensorscalar(C)Return the single element of a tensor-like object with zero indices or dimensions as a value of the underlying scalar type.
Allocations
For some networks, it will be necessary to allocate output and/or intermediate tensors. This process is split into the following hierarchy of functions, where custom tensor types can opt in at different levels.
By default, the process is split into three steps.
First, the scalar type TC of the resulting tensor is determined. This is done by leveraging VectorInterface.jl's scalartype, and promoting the results along with the types of any scalars that are present.
TensorOperations.promote_add — Functionpromote_add(args...)Promote the scalar types of a tensor addition to a common type.
TensorOperations.promote_contract — Functionpromote_contract(args...)Promote the scalar types of a tensor contraction to a common type.
Then, the type and structure of the resulting tensor is determined. The former represents all the information that is contained within the type, while the latter adds the required runtime information (e.g. array sizes, ...).
TensorOperations.tensoradd_type — Functiontensoradd_type(TC, pC, A, conjA)Obtain typeof(C), where C is the result of tensoradd!(C, pC, A, conjA) with scalar type TC.
TensorOperations.tensoradd_structure — Functiontensoradd_structure(pC, A, conjA)Obtain the structure information of C, where C would be the output of tensoradd!(C, pC, A, conjA).
TensorOperations.tensorcontract_type — Functiontensorcontract_type(TC, pC, A, pA, conjA, B, pB, conjB)Obtain typeof(C), where C is the result of tensorcontract!(C, pC, A, pA, conjA, B, pB, conjB) with scalar type TC.
TensorOperations.tensorcontract_structure — Functiontensorcontract_structure(pC, A, pA, conjA, B, pB, conjB)Obtain the structure information of C, where C would be the output of tensorcontract!(C, pC, A, pA, conjA, B, pB, conjB).
Finally, the tensor is allocated, where a flag indicates if this is a temporary object, or one that will persist outside of the scope of the macro. If the resulting tensor is a temporary object and its memory will not be freed by Julia's garbage collector, it can be explicitly freed by implementing tensorfree!, which by default does nothing.
TensorOperations.tensoralloc — Functiontensoralloc(ttype, structure, istemp=false, [backend::Backend])Allocate memory for a tensor of type ttype and structure structure. The optional third argument can be used to indicate that the result is used as a temporary tensor, for which in some cases and with some backends (the optional fourth argument) a different allocation strategy might be used.
See also tensoralloc_add, tensoralloc_contract and tensorfree!.
TensorOperations.tensorfree! — Functiontensorfree!(C, [backend::Backend])Provide a hint that the allocated memory of C can be released.
See also tensoralloc.
The @tensor macro will however only insert the calls to the following functions, which have a default implementation in terms of the functions above.
TensorOperations.tensoralloc_add — Functiontensoralloc_add(TC, pC, A, conjA, istemp=false, backend::Backend...)Allocate a tensor C of scalar type TC that would be the result of
`tensoradd!(C, pC, A, conjA)`The istemp argument is used to indicate that a tensor wlil not be used after the @tensor block, and thus will be followed by an explicit call to tensorfree!. The backend can be used to implement different allocation strategies.
See also tensoralloc and tensorfree!.
TensorOperations.tensoralloc_contract — Functiontensoralloc_contract(TC, pC, A, pA, conjA, B, pB, conjB, istemp=false, backend::Backend...)Allocate a tensor C of scalar type TC that would be the result of
`tensorcontract!(C, pC, A, pA, conjA, B, pB, conjB)`The istemp argument is used to indicate that a tensor wlil not be used after the @tensor block, and thus will be followed by an explicit call to tensorfree!. The backend can be used to implement different allocation strategies.
See also tensoralloc and tensorfree!.
Utility
Some of the optional keywords for @tensor can be accessed only after implementing the following utility functions:
TensorOperations.tensorcost — Functiontensorcost(A, i)Compute the contraction cost associated with the ith index of a tensor, such that the total cost of a pairwise contraction is found as the product of the costs of all contracted indices and all uncontracted indices.
TensorOperations.checkcontractible — Functioncheckcontractible(A, iA, conjA, B, iB, conjB, label)Verify whether two tensors opA(A) and opB(B) are compatible for having their respective index iA and iB contracted, and throws an error if not. The operation opA acts as identity if conjA equals :N and as conj if conjA equals :C; the operation opB is determined by conjB analogously.