Key Features

• Extensive mixed-precision support:

  • FP64 inputs with FP32 compute.

  • FP32 inputs with FP16, BF16, or TF32 compute.

  • Complex-times-real operations.

  • Conjugate (without transpose) support.

• Support for up to 64-dimensional tensors.

• Arbitrary data layouts.

• Trivially serializable data structures.

• Multi-process distributed tensor contractions with cuTENSORMp (Linux only).

• Main computational routines:

  • Direct (i.e., transpose-free) tensor contractions.

    • Support just-in-time compilation of dedicated kernels.

  • Tensor reductions (including partial reductions).

  • Element-wise tensor operations:

    • Support for various activation functions.

    • Support for padding of the output tensor

    • Arbitrary tensor permutations.

    • Conversion between different data types.

Documentation

Please refer to https://docs.nvidia.com/cuda/cutensor/index.html for the cuTENSOR documentation.

Installation

The cuTENSOR wheel can be installed as follows:

pip install cutensor-cuXX

where XX is the CUDA major version (currently CUDA 12 & 13 are supported). The package cutensor (without the -cuXX suffix) is deprecated. If you have cutensor installed, please remove it prior to installing cutensor-cuXX.

On Linux, install the optional mp extra to use cuTENSORMp. This installs the matching NVIDIA NCCL runtime:

pip install "cutensor-cuXX[mp]"

The mp extra installs NCCL only. The NCCL-backed cuTENSORMp library is not linked against MPI, so the wheel deliberately does not select an MPI implementation. MPI-based applications and samples require an external MPI development installation and launcher, such as Open MPI, MPICH, or NVIDIA HPC-X. Compiling the samples also requires a CUDA development toolkit with nvcc.