https://github.com/jcmgray/einsum_bmm

einsum via batch matrix multiply
https://github.com/jcmgray/einsum_bmm

einsum tensor tensor-contraction tensor-networks

Last synced: about 2 months ago
JSON representation

einsum via batch matrix multiply

• Host: GitHub
• URL: https://github.com/jcmgray/einsum_bmm
• Owner: jcmgray
• License: gpl-3.0
• Created: 2023-02-16T00:25:09.000Z (almost 2 years ago)
• Default Branch: main
• Last Pushed: 2023-11-29T19:09:43.000Z (about 1 year ago)
• Last Synced: 2024-10-31T08:52:18.490Z (2 months ago)
• Topics: einsum, tensor, tensor-contraction, tensor-networks
• Language: Python
• Homepage:
• Size: 37.1 KB
• Stars: 13
• Watchers: 2
• Forks: 0
• Open Issues: 0
• Metadata Files:
  • Readme: README.md
  • License: LICENSE

Awesome Lists containing this project

README

        
# `einsum_bmm`

**:cactus: This has now been incorporated into `cotengra`, see the [high level functionality](https://cotengra.readthedocs.io/en/latest/high-level-interface.html), called with `implementation='cotengra'`, and the actual two term implementation [here](https://cotengra.readthedocs.io/en/latest/autoapi/cotengra/contract/index.html#cotengra.contract.einsum). There is also a PR for `numpy.einsum` with the `optimize` kwarg [here](https://github.com/numpy/numpy/pull/23513) :cactus:**

This repository provides an `einsum` (and `tensordot`) function implemented via **batch matrix

multiply**.

1. This *can* be much faster than the raw `numpy.einsum` function, especially

   for large and high dimensional contractions.

2. It can also be used to enable `einsum` for any backend that provides only

   `tranpose`, `reshape` and `matmul`.

The implementation is achieved by grouping indices according to the following classification:



1. Summed indices are trivially removed.

2. A and B and then transposed and reshaped for batched matrix multiplication

3. The output is reshaped and transposed

Each of these steps only occurs if necessary. There are slight specializations for both pure multiplication and no batch indices.

Notes:

* It currently only supports 1 or 2 terms, a library such as `opt_einsum` or

  `cotengra` should be used to dispatch many term contractions to a pairwise

  ordering in conjuction with this `einsum_bmm`.