A symmetric tensor
These arise in a variety of fields within chemistry and physics, where the symmetry is prescribed by the physics. The physics often also prescribe that the tensor be high dimensional and/or high rank, which can quickly exceed memory capacity if the tensor is stored in a normal NumPy array. Thankfully, symmetry also presents a solution to this problem: since many of the tensor components are redundant, they can be stored in a compressed format. However, it is not sufficient to simply store these tensors: we also need to use them. Therefore the interest in wrapping their compressed storage with a NumPy-compatible API, exposing standard slicing and arithmetic semantics, allowing mixture and replacement with standard NumPy arrays, and completely hiding the compressed memory layout.1 Crucially, whenever possible, operations are performed directly on the underlying compressed memory, which in some case can lead order-of-magnitude gains in efficiency – or even make the impossible, possible.2
This package evolved out of a desire to apply field-theory methods to neural networks, which required tensors of both high rank and very high dimension to represent polynomials. Since evaluating these polynomials implicitely symmetrizes all the tensors (for example,
This package therefore addresses three complementary needs:
- Provide ready-to-use implementations of symmetric tensors, which behave like standard arrays wherever possible.
- Add support for a selection of symmetrized linear algebra operations, which exploit the symmetric structure of tensors for sometimes massive computational gains. Generic fallback operations are provided for non-symmetric arrays.
- Make it easy to extend support for:
- additional functions;
- additional array backends (currently NumPy and PyTorch are suported)
- new storage compressed storage formats.
Just like with sparse arrays, different applications may be best served by different compressed storage formats. This is possible by subclassing the base SymmetricTensor class and implementing format-specific methods. The resulting class can be made completely interoperable with NumPy arrays and other SymmetricTensor classes.
In addition to compression formats, one may also want to use different backends than NumPy for the data arrays, such as PyTorch or JaX. This is done via mixin classes. By combining mixins with SymmetricTensor subclasses, a wide array of tensor formats can be supported with minimal repeated code.
The functions below define symmetrized versions of the equivalent NumPy function. In other words, if a NumPy function can be written
where
Since symmetrized functions are closed over symmetric tensors, we can say that they form a symmetric algebra; they are correspondingly defined in symtensor.symalg. Note that you should always use the functions defined in this module: they will automatically dispatch to an optimized implementation when one is available.
Currently symtensor.symalg defines the following functions:
addaddadd.outer
subtractsubtractsubtract.outer
multiplymultiplymultiply.outer
transpose(optimized to a no-op on symmetric tensors)tensordotcontract_all_indices_with_matrixcontract_all_indices_with_vectorcontract_tensor_list
| Format → Backend ↓ |
Dense | Permutation classes | Outer product decomposition | Blocked arrays |
|---|---|---|---|---|
| Numpy | ✔ | ✔ | (WIP) | ✘ |
| PyTorch | ✔ | ✔ | prototype | ✘ |
| tlash3 | ✘ | ✘ | ✘ | (planned) |
We use a standardized API suite of unit tests: most tests are written generically enough to be applied to each format/backend implementation without any modification. This ensures a consistent API throughout the library, and massively simplifies the writing of tests for new implementations. In brief, porting the test suite involves:
- Copying one of the existing tests under tests.
- Opening the file as a notebook via Jupytext. This provides a lot of inlined documentation, and organizational sections to aid navigation.
- Changing the few lines at the top which specify the subclass of SymmetricTensor to test.
- Adding/deactivating/modifying tests as needed for the particulars of your format and backend. Typically, out of about 25 tests, 3-4 need to be specialized.
Footnotes
-
We implement both the universal function and array function protocols. These are also called dispatch mechanisms. ↩
-
These requirements are very similar to those for sparse arrays, and this package makes extensive use of NumPy improvements introduced to address limitations of scipy.sparse. In fact our base class could serve as a basis for a new implementation of sparse arrays, if one were the type inclined to reinvent the wheel. ↩
-
tlashis a C implementation of the blocked arrays format described by Schatz et al. (Exploiting Symmetry in Tensors for High Performance: Multiplication with Symmetric Tensors, 2014). [GitHub] [doi:10.1137/130907215] ↩