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 boolean values that are used to indicate if the input tensor should be conjugated (true) or used as-is (false).
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 and/or allocator. For new types of tensors, it is possible to implement the methods without the final two arguments backend and allocator, if you know that you will never need them. They will not be inserted by the @tensor macro if they are not explicitly specified as keyword arguments. However, it is possible to add those arguments with the default values and ignore them in case they are not needed.
Alternatively, if some new tensor type is backed by an AbstractArray instance, and the tensor operations are also implemented by applying the same operations to the underlying array, it is possible to forward the value of the backend and allocator arguments in order to still support the full flexibility of the @tensor macro.
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.
As there is typically a simple and unique way of doing so, this method does not have any backend or allocator arguments.
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.tensorstructure — Functiontensorstructure(A, iA, conjA)Obtain the information associated to indices iA of tensor op(A), where op acts as identity if conjA equals false and as conj if conjA equals true.
TensorOperations.tensoradd_type — Functiontensoradd_type(TC, A, pA, conjA)Obtain the type information of C, where C would be the output of tensoradd!(C, A, pA, conjA) with scalartype TC.
TensorOperations.tensoradd_structure — Functiontensoradd_structure(A, pA, conjA)Obtain the structure information of C, where C would be the output of tensoradd!(C, A, pA, conjA).
TensorOperations.tensorcontract_type — Functiontensorcontract_type(TC, A, pA, conjA, B, pB, conjB, pAB)Obtain the type information of C, where C would be the output of tensorcontract!(C, A, pA, conjA, B, pB, conjB, pAB) with scalar type TC.
TensorOperations.tensorcontract_structure — Functiontensorcontract_structure(A, pA, conjA, B, pB, conjB, pAB)Obtain the structure information of C, where C would be the output of tensorcontract!(C, A, pA, conjA, B, pB, conjB, pAB).
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, allocator])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 allocators (the optional fourth argument) a different allocation strategy might be used.
See also tensoralloc_add, tensoralloc_contract and tensorfree!.
TensorOperations.tensorfree! — Functiontensorfree!(C, [allocator])Provide a hint that the allocated memory of C can be released.
See also tensoralloc.
These functions also depend on the optional allocator argument that can be used to control different allocation strategies, and for example to differentiate between allocations of temporary tensors as opposed to tensors that are part of the output.
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, A, pA, conjA, [istemp=Val(false), allocator])Allocate a tensor C of scalar type TC that would be the result of
`tensoradd!(C, A, pA, 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 allocator can be used to implement different allocation strategies.
See also tensoralloc and tensorfree!.
TensorOperations.tensoralloc_contract — Functiontensoralloc_contract(TC, A, pA, conjA, B, pB, conjB, pAB, [istemp=Val(false), allocator])Allocate a tensor C of scalar type TC that would be the result of
`tensorcontract!(C, A, pA, conjA, B, pB, conjB, pAB)`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 allocator 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 false and as conj if conjA equals true; the operation opB is determined by conjB analogously.