pytblis 0.0.9

Python bindings for TBLIS, a high performance library for tensor operations

pytblis: Python bindings for TBLIS

Are your einsums too slow?

Need FP64 tensor contractions and can't buy a datacenter GPU because you already maxed out your home equity line of credit?

Set your CPU on fire with TBLIS!

Installation

pip install pytblis

The pre-built wheels on PyPI use pthreads for multithreading. To reduce the overhead due to creating and joining threads, compile pytblis yourself and configure it to use OpenMP.

Building

1. Install TBLIS.
2. Run CMAKE_ARGS="-DTBLIS_ROOT=wherever_tblis_is_installed" pip install .

See dev_install.sh for an example. This script installs TBLIS in ./local_tblis_prefix and then links pytblis against it.

Usage

pytblis.einsum and pytblis.tensordot are drop-in replacements for numpy.einsum and numpy.tensordot.

In addition, low level wrappers are provided for tblis_tensor_add, tblis_tensor_mult, tblis_tensor_reduce, tblis_tensor_shift, and tblis_tensor_dot. These are named pytblis.add, pytblis.mult, et cetera.

Finally, there are mid-level convenience wrappers for tblis_tensor_mult and tblis_tensor_add:

def contract(
    subscripts: str,
    a: ArrayLike,
    b: ArrayLike,
    alpha: scalar = 1.0,
    beta: scalar = 0.0,
    out: Optional[npt.ArrayLike] = None,
    conja: bool = False,
    conjb: bool = False,
) -> ArrayLike

and

def transpose_add(
    subscripts: str,
    a: ArrayLike,
    alpha: scalar = 1.0,
    beta: scalar = 0.0,
    out: Optional[ArrayLike] = None,
    conja: bool = False,
    conjout: bool = False,
) -> ArrayLike

These are used as follows:

C = pytblis.contract("ij,jk->ik", A, B, alpha=1.0, beta=0.5, out=C, conja=True, conjb=False)

does

$$C \gets \overline{A} B + \frac{1}{2} C.$$

B = pytblis.tensor_add("iklj->ijkl", A, alpha=-1.0, beta=1.0, out=B)

does

$$B_{ijkl} \gets B_{ijkl} - A_{iklj}.$$

Some additional documentation (work in progress) is available at pytblis.readthedocs.io.

Limitations

Supported datatypes: np.float32, np.float64, np.complex64, np.complex128. Mixing arrays of different types isn't yet supported. I may add a workaround for real-complex tensor contraction.

Arrays with negative or zero stride are not supported and will cause pytblis to fall back to NumPy (for einsum and contract) or raise an error (all other functions).

Research

If you use TBLIS in your academic work, it's a good idea to cite:

• High-Performance Tensor Contraction without Transposition
• Strassen's Algorithm for Tensor Contraction

TBLIS is not my work, and its developers are not responsible for flaws in these Python bindings.

Acknowledgements

The implementation of einsum and the tests are modified versions of those from opt_einsum.