Tensor Operations#

Overview#

The tensor module nvmath.tensor in nvmath-python provides APIs for tensor operations powered by the high-performance NVIDIA cuTENSOR library. We currently offer binary and ternary contraction APIs supporting the CUDA execution space.

For contracting a tensor network, refer to the Network API from the cuQuantum library. While network contraction can be used for binary and ternary contraction, the focus here is on the optimal contraction of a tensor network and therefore not all options pertinent to each pairwise contraction are available to the user. The generalized binary \(\alpha \; a \cdot b + \beta \; c\) and ternary \(\alpha \; a \cdot b \cdot c + \beta \; d\) contraction operations (where \(\cdot\) represents tensor contraction) in this module are fused, and support options specific to efficient execution of these operations.

import cupy as cp
from cupyx.profiler import benchmark

import nvmath

a = cp.random.rand(64, 8, 8, 6, 6)
b = cp.random.rand(64, 8, 8, 6, 6)

# Create a stateful BinaryContraction object 'contraction'.
with nvmath.tensor.BinaryContraction("pijkl,pjiab->lakbp", a, b) as contraction:
    # Get the handle to the plan preference object
    plan_preference = contraction.plan_preference
    # update the kernel rank to the third best for the underlying algorithm
    plan_preference.kernel_rank = 2

    for algo in (
        nvmath.tensor.ContractionAlgo.DEFAULT_PATIENT,
        nvmath.tensor.ContractionAlgo.GETT,
        nvmath.tensor.ContractionAlgo.TGETT,
        nvmath.tensor.ContractionAlgo.TTGT,
        nvmath.tensor.ContractionAlgo.DEFAULT,
    ):
        print(f"Algorithm: {algo.name}")
        plan_preference.algo = algo
        # Plan the Contraction to activate the updated plan preference
        contraction.plan()
        print(benchmark(contraction.execute, n_repeat=20))

More examples of tensor operations can be found on our GitHub repository.

Host API Reference#