Skip to main content

NVIDIA cuTENSOR

Project description

cuTENSOR is a high-performance CUDA library for tensor primitives.

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.

  • Main computational routines:

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

    • Tensor reductions (including partial reductions).

    • Element-wise tensor operations:

      • Support for various activation functions.

      • 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 11 & 12 are supported). The package cutensor (without the -cuXX suffix) is considered deprecated. If you have cutensor installed, please remove it prior to installing cutensor-cuXX.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page