pytblis 0.0.4

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

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.