cuquantum-python-jax-cu12 0.0.5__tar.gz

cuquantum_python_jax_cu12-0.0.5/LICENSE

@@ -0,0 +1,28 @@
+Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+BSD-3-Clause
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+this list of conditions and the following disclaimer in the documentation
+and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+contributors may be used to endorse or promote products derived from
+this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

cuquantum_python_jax_cu12-0.0.5/MANIFEST.in

@@ -0,0 +1,9 @@
+include pyproject.toml.template
+include configure.sh
+include cuquantum/densitymat/jax/cppsrc/CMakeLists.txt
+include cuquantum/densitymat/jax/cppsrc/cudensitymat.h
+include cuquantum/densitymat/jax/cppsrc/cudensitymat_jax.cpp
+include cuquantum/densitymat/jax/cppsrc/cudensitymat_jax.h
+include cuquantum/densitymat/jax/cppsrc/pybind.cpp
+include cuquantum/densitymat/jax/cppsrc/utils.h
+prune tests*

cuquantum_python_jax_cu12-0.0.5/PKG-INFO

@@ -0,0 +1,113 @@
+Metadata-Version: 2.4
+Name: cuquantum-python-jax-cu12
+Version: 0.0.5
+Summary: NVIDIA cuQuantum Python JAX
+Author-email: NVIDIA Corporation <cuquantum-python@nvidia.com>
+License-Expression: BSD-3-Clause
+Project-URL: Homepage, https://developer.nvidia.com/cuquantum-sdk
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Topic :: Education
+Classifier: Topic :: Scientific/Engineering
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Environment :: GPU :: NVIDIA CUDA
+Classifier: Environment :: GPU :: NVIDIA CUDA :: 12
+Requires-Python: >=3.11.0
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: pybind11
+Requires-Dist: cuquantum-python-cu12~=26.1.0
+Requires-Dist: jax[cuda12-local]<0.7,>=0.5
+Dynamic: license-file
+
+# cuQuantum Python JAX
+
+cuQuantum Python JAX provides a JAX extension for cuQuantum Python. It exposes selected functionality of cuQuantum SDK in a JAX-compatible way that enables JAX frameworks to directly interface with the exposed cuQuantum API. In the current release, cuQuantum JAX exposes a JAX interface to the Operator Action API from the cuDensityMat library.
+
+## Documentation
+
+Please visit the [NVIDIA cuQuantum Python documentation](https://docs.nvidia.com/cuda/cuquantum/latest/python).
+
+## Building and installing cuQuantum Python JAX
+
+### Requirements
+
+The install-time dependencies of the cuQuantum Python JAX package include:
+
+* cuquantum-python-cu12~=26.1.0 for CUDA 12 or cuquantum-python-cu13~=26.1.0 for CUDA 13
+* jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8,<0.9 for CUDA 13
+* pybind11
+* setuptools>=77.0.3
+
+Note:
+1. cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13.
+2. cuQuantum Python JAX installation does not support build isolation. The user needs to pass in `--no-build-isolation` to `pip` when installing cuQuantum Python JAX.
+3. cuQuantum Python JAX wheels are CUDA-versioned: `cuquantum-python-jax-cu12` for CUDA 12 and `cuquantum-python-jax-cu13` for CUDA 13.
+
+#### Installation using `jax[cudaXX-local]`
+
+`cuquantum-python-jax-cu12` (or `cuquantum-python-jax-cu13`) depends explicitly on `jax[cudaXX-local]`. Installing the package will also install `jax[cudaXX-local]`.
+
+Using `jax[cudaXX-local]` assumes the user provides both cuDNN and the CUDA Toolkit. cuDNN is not a part of the CUDA Toolkit and requires an additional installation. The user must also specify `LD_LIBRARY_PATH`, including the library folders containing `libcudnn.so` and `libcupti.so`.
+
+`libcupti.so` is provided by the CUDA Toolkit. If the CUDA Toolkit is installed under `/usr/local/cuda`, `libcupti.so` is located under `/usr/local/cuda/extras/CUPTI/lib64` and `LD_LIBRARY_PATH` should contain this path.
+
+`libcudnn.so` is installed separately from the CUDA Toolkit. The default installation location is `/usr/local/cuda/lib64`, and `LD_LIBRARY_PATH` should contain this path.
+
+Both `libcudnn.so` and `libcupti.so` are installable with pip:
+
+```
+pip install nvidia-cudnn-cu12
+pip install nvidia-cuda-cupti-cu12
+```
+
+After installing cuDNN and cuPTI, the user may install cuQuantum Python JAX with `pip` using either:
+
+```
+pip install --no-build-isolation cuquantum-python-jax-cu12   # for CUDA 12
+pip install --no-build-isolation cuquantum-python-jax-cu13   # for CUDA 13
+```
+
+or one of
+
+```
+pip install --no-build-isolation cuquantum-python-cu12[jax]
+pip install --no-build-isolation cuquantum-python-cu13[jax]
+```
+
+where the CUDA version is explicitly specified on cuquantum-python.
+
+Note:
+1. If cuDNN and cuPTI are installed with `pip`, the user does not need to specify library folders in `LD_LIBRARY_PATH`.
+2. When the latter command `pip install --no-build-isolation cuquantum-python-cu12[jax]`/`pip install --no-build-isolation cuquantum-python-cu13[jax]` is used, `--no-build-isolation` applies to both cuquantum-python and cuquantum-python-jax. The user needs to ensure cuquantum-python's build dependencies are installed before the installation.
+
+#### Installing from source
+
+To install cuQuantum Python JAX from source, first compile cuQuantum Python from source using the [instructions on GitHub](https://github.com/NVIDIA/cuQuantum/blob/main/python/README.md). Once complete, navigate to `python/extensions`, then:
+
+```
+pip install .
+```
+
+The CUDA version is detected automatically from `$CUDA_PATH` and the wheel will be named accordingly (`cuquantum-python-jax-cu12` or `cuquantum-python-jax-cu13`).
+
+## Running
+
+### Requirements
+
+Runtime dependencies of the cuQuantum Python JAX package include:
+
+* An NVIDIA GPU with compute capability 7.5+
+* cuquantum-python-cu12~=26.1.0 for CUDA 12 or cuquantum-python-cu13~=26.1.0 for CUDA 13
+* jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8,<0.9 for CUDA 13
+* pybind11
+
+## Developer Notes
+
+* cuQuantum Python JAX does not support editable installation.
+* Both cuQuantum Python and cuQuantum Python JAX need to be installed into `site-packages` for proper import of the library.
+* cuQuantum Python JAX assumes cuQuantum Python will be available under the current `site-packages` directory.

cuquantum_python_jax_cu12-0.0.5/README.md

@@ -0,0 +1,87 @@
+# cuQuantum Python JAX
+
+cuQuantum Python JAX provides a JAX extension for cuQuantum Python. It exposes selected functionality of cuQuantum SDK in a JAX-compatible way that enables JAX frameworks to directly interface with the exposed cuQuantum API. In the current release, cuQuantum JAX exposes a JAX interface to the Operator Action API from the cuDensityMat library.
+
+## Documentation
+
+Please visit the [NVIDIA cuQuantum Python documentation](https://docs.nvidia.com/cuda/cuquantum/latest/python).
+
+## Building and installing cuQuantum Python JAX
+
+### Requirements
+
+The install-time dependencies of the cuQuantum Python JAX package include:
+
+* cuquantum-python-cu12~=26.1.0 for CUDA 12 or cuquantum-python-cu13~=26.1.0 for CUDA 13
+* jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8,<0.9 for CUDA 13
+* pybind11
+* setuptools>=77.0.3
+
+Note:
+1. cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13.
+2. cuQuantum Python JAX installation does not support build isolation. The user needs to pass in `--no-build-isolation` to `pip` when installing cuQuantum Python JAX.
+3. cuQuantum Python JAX wheels are CUDA-versioned: `cuquantum-python-jax-cu12` for CUDA 12 and `cuquantum-python-jax-cu13` for CUDA 13.
+
+#### Installation using `jax[cudaXX-local]`
+
+`cuquantum-python-jax-cu12` (or `cuquantum-python-jax-cu13`) depends explicitly on `jax[cudaXX-local]`. Installing the package will also install `jax[cudaXX-local]`.
+
+Using `jax[cudaXX-local]` assumes the user provides both cuDNN and the CUDA Toolkit. cuDNN is not a part of the CUDA Toolkit and requires an additional installation. The user must also specify `LD_LIBRARY_PATH`, including the library folders containing `libcudnn.so` and `libcupti.so`.
+
+`libcupti.so` is provided by the CUDA Toolkit. If the CUDA Toolkit is installed under `/usr/local/cuda`, `libcupti.so` is located under `/usr/local/cuda/extras/CUPTI/lib64` and `LD_LIBRARY_PATH` should contain this path.
+
+`libcudnn.so` is installed separately from the CUDA Toolkit. The default installation location is `/usr/local/cuda/lib64`, and `LD_LIBRARY_PATH` should contain this path.
+
+Both `libcudnn.so` and `libcupti.so` are installable with pip:
+
+```
+pip install nvidia-cudnn-cu12
+pip install nvidia-cuda-cupti-cu12
+```
+
+After installing cuDNN and cuPTI, the user may install cuQuantum Python JAX with `pip` using either:
+
+```
+pip install --no-build-isolation cuquantum-python-jax-cu12   # for CUDA 12
+pip install --no-build-isolation cuquantum-python-jax-cu13   # for CUDA 13
+```
+
+or one of
+
+```
+pip install --no-build-isolation cuquantum-python-cu12[jax]
+pip install --no-build-isolation cuquantum-python-cu13[jax]
+```
+
+where the CUDA version is explicitly specified on cuquantum-python.
+
+Note:
+1. If cuDNN and cuPTI are installed with `pip`, the user does not need to specify library folders in `LD_LIBRARY_PATH`.
+2. When the latter command `pip install --no-build-isolation cuquantum-python-cu12[jax]`/`pip install --no-build-isolation cuquantum-python-cu13[jax]` is used, `--no-build-isolation` applies to both cuquantum-python and cuquantum-python-jax. The user needs to ensure cuquantum-python's build dependencies are installed before the installation.
+
+#### Installing from source
+
+To install cuQuantum Python JAX from source, first compile cuQuantum Python from source using the [instructions on GitHub](https://github.com/NVIDIA/cuQuantum/blob/main/python/README.md). Once complete, navigate to `python/extensions`, then:
+
+```
+pip install .
+```
+
+The CUDA version is detected automatically from `$CUDA_PATH` and the wheel will be named accordingly (`cuquantum-python-jax-cu12` or `cuquantum-python-jax-cu13`).
+
+## Running
+
+### Requirements
+
+Runtime dependencies of the cuQuantum Python JAX package include:
+
+* An NVIDIA GPU with compute capability 7.5+
+* cuquantum-python-cu12~=26.1.0 for CUDA 12 or cuquantum-python-cu13~=26.1.0 for CUDA 13
+* jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8,<0.9 for CUDA 13
+* pybind11
+
+## Developer Notes
+
+* cuQuantum Python JAX does not support editable installation.
+* Both cuQuantum Python and cuQuantum Python JAX need to be installed into `site-packages` for proper import of the library.
+* cuQuantum Python JAX assumes cuQuantum Python will be available under the current `site-packages` directory.

cuquantum_python_jax_cu12-0.0.5/configure.sh

@@ -0,0 +1,122 @@
+#!/usr/bin/env bash
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Generate pyproject.toml from pyproject.toml.template by detecting
+# CUDA version and substituting template placeholders.
+#
+# Usage:
+#   CUDA_PATH=/usr/local/cuda bash configure.sh
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+TEMPLATE="${SCRIPT_DIR}/pyproject.toml.template"
+OUTPUT="${SCRIPT_DIR}/pyproject.toml"
+
+# -------------------------------------------------------------------
+# 1. Validate template exists (early, before doing any real work)
+# -------------------------------------------------------------------
+if [[ ! -f "${TEMPLATE}" ]]; then
+    echo "ERROR: Template not found at ${TEMPLATE}" >&2
+    exit 1
+fi
+
+# -------------------------------------------------------------------
+# 2. Validate CUDA_PATH
+# -------------------------------------------------------------------
+if [[ -z "${CUDA_PATH:-}" ]]; then
+    echo "ERROR: CUDA_PATH is not set. Please set it to your CUDA toolkit root." >&2
+    exit 1
+fi
+
+CUDA_H="${CUDA_PATH}/include/cuda.h"
+if [[ ! -f "${CUDA_H}" ]]; then
+    echo "ERROR: Cannot find ${CUDA_H}. Is CUDA_PATH set correctly?" >&2
+    exit 1
+fi
+
+# -------------------------------------------------------------------
+# 3. Parse CUDA_VERSION from cuda.h  (mirrors setup.py logic)
+#    Example line:  #define CUDA_VERSION 12020
+#    12020 => major = 12020 / 1000 = 12
+# -------------------------------------------------------------------
+# Read grep output into a variable to avoid a grep|head|sed pipeline,
+# which can behave inconsistently with pipefail across bash versions
+# (head closes the pipe early, causing grep to receive SIGPIPE).
+CUDA_H_MATCH=$(grep -E -m1 '^\s*#define\s+CUDA_VERSION\s+[0-9]+' "${CUDA_H}" || true)
+
+if [[ -z "${CUDA_H_MATCH}" ]]; then
+    echo "ERROR: Could not find CUDA_VERSION definition in ${CUDA_H}" >&2
+    exit 1
+fi
+
+CUDA_VERSION_RAW=$(echo "${CUDA_H_MATCH}" | sed -E 's/.*#define\s+CUDA_VERSION\s+([0-9]+).*/\1/')
+
+if [[ -z "${CUDA_VERSION_RAW}" ]]; then
+    echo "ERROR: Could not parse CUDA_VERSION from ${CUDA_H}" >&2
+    exit 1
+fi
+
+if (( CUDA_VERSION_RAW < 1000 )); then
+    echo "ERROR: CUDA_VERSION ${CUDA_VERSION_RAW} from ${CUDA_H} is unexpectedly small (< 1000)" >&2
+    exit 1
+fi
+
+CUDA_MAJOR=$(( CUDA_VERSION_RAW / 1000 ))
+
+echo "Detected CUDA_VERSION=${CUDA_VERSION_RAW} (major=${CUDA_MAJOR}) from ${CUDA_H}"
+
+# -------------------------------------------------------------------
+# 4. Map CUDA major version to template variable values
+# -------------------------------------------------------------------
+case "${CUDA_MAJOR}" in
+    12)
+        JAX_VERSION_SPEC=">=0.5,<0.7"
+        CUDA_CLASSIFIER="Environment :: GPU :: NVIDIA CUDA :: 12"
+        ;;
+    13)
+        JAX_VERSION_SPEC=">=0.8,<0.9"
+        CUDA_CLASSIFIER="Environment :: GPU :: NVIDIA CUDA :: 13"
+        ;;
+    *)
+        echo "ERROR: Unsupported CUDA major version: ${CUDA_MAJOR}" >&2
+        exit 1
+        ;;
+esac
+
+# -------------------------------------------------------------------
+# 5. Escape sed replacement-special characters (& and \) in values
+# -------------------------------------------------------------------
+sed_escape() {
+    printf '%s' "$1" | sed -e 's/[&\\/]/\\&/g'
+}
+
+CUDA_MAJOR_ESC=$(sed_escape "${CUDA_MAJOR}")
+JAX_VERSION_SPEC_ESC=$(sed_escape "${JAX_VERSION_SPEC}")
+CUDA_CLASSIFIER_ESC=$(sed_escape "${CUDA_CLASSIFIER}")
+
+# -------------------------------------------------------------------
+# 6. Generate pyproject.toml via sed substitution
+#    Write to a temp file then atomically move into place so a partial
+#    write (e.g. Ctrl+C, disk full) never leaves a broken pyproject.toml.
+# -------------------------------------------------------------------
+TMPFILE=$(mktemp "${OUTPUT}.XXXXXX")
+trap 'rm -f "${TMPFILE}"' EXIT
+
+sed -e "s|@CUDA_MAJOR_VER@|${CUDA_MAJOR_ESC}|g" \
+    -e "s|@JAX_VERSION_SPEC@|${JAX_VERSION_SPEC_ESC}|g" \
+    -e "s|@CUDA_CLASSIFIER@|${CUDA_CLASSIFIER_ESC}|g" \
+    -- "${TEMPLATE}" > "${TMPFILE}"
+
+# Skip overwrite if output already exists and is identical (idempotency).
+if [[ -f "${OUTPUT}" ]] && cmp -s "${TMPFILE}" "${OUTPUT}"; then
+    echo "pyproject.toml is already up to date for CUDA ${CUDA_MAJOR}"
+    exit 0
+fi
+
+mv -f "${TMPFILE}" "${OUTPUT}"
+trap - EXIT
+
+echo "Generated ${OUTPUT} for CUDA ${CUDA_MAJOR}"

cuquantum_python_jax_cu12-0.0.5/cuquantum/densitymat/jax/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from cuquantum.bindings._internal import cudensitymat as _cudm
+_cudm._inspect_function_pointers()  # for loading libcudensitymat.so
+
+import jax
+if not jax.config.jax_enable_x64:
+    raise RuntimeError(f"jax_enable_x64 must be set to True to use cuQuantum Python JAX")
+
+from .operator_action import operator_action
+from .pysrc import (
+    ElementaryOperator,
+    MatrixOperator,
+    OperatorTerm,
+    Operator
+)

cuquantum_python_jax_cu12-0.0.5/cuquantum/densitymat/jax/cppsrc/CMakeLists.txt

@@ -0,0 +1,68 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+cmake_minimum_required(VERSION 3.22)
+project(cudensitymat_jax LANGUAGES CXX)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+find_package(Python3 REQUIRED COMPONENTS Interpreter Development)
+message(STATUS "Python executable: ${Python3_EXECUTABLE}")
+
+find_package(CUDAToolkit REQUIRED)
+message(STATUS "CUDA toolkit directory: ${CUDAToolkit_INCLUDE_DIRS}")
+
+# Find XLA directory
+execute_process(
+    COMMAND ${Python3_EXECUTABLE} -c "import jax; print(jax.ffi.include_dir())"
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    OUTPUT_VARIABLE XLA_DIR
+)
+if(NOT XLA_DIR)
+    message(FATAL_ERROR "XLA directory not found")
+else()
+    message(STATUS "XLA directory: ${XLA_DIR}")
+endif()
+
+# Find pybind11 directory
+execute_process(
+    COMMAND ${Python3_EXECUTABLE} -c "import pybind11; print(pybind11.get_include())"
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    OUTPUT_VARIABLE pybind11_INCLUDE_DIR
+)
+if(NOT pybind11_INCLUDE_DIR)
+    message(FATAL_ERROR "Pybind11 include directory not found")
+else()
+    message(STATUS "Pybind11 include directory: ${pybind11_INCLUDE_DIR}")
+endif()
+
+set(pybind11_DIR ${pybind11_INCLUDE_DIR}/../share/cmake/pybind11)
+find_package(pybind11 REQUIRED)
+
+pybind11_add_module(
+    ${PROJECT_NAME}
+    cudensitymat_jax.cpp
+    pybind.cpp
+)
+target_include_directories(
+    ${PROJECT_NAME}
+    PUBLIC
+    ${CUDAToolkit_INCLUDE_DIRS}
+    ${XLA_DIR}
+    ${pybind11_INCLUDE_DIR}
+    ${CMAKE_CURRENT_SOURCE_DIR}
+)
+
+set_target_properties(
+    ${PROJECT_NAME}
+    PROPERTIES
+        BUILD_RPATH "$ORIGIN"
+)
+
+target_link_libraries(
+    ${PROJECT_NAME}
+    PRIVATE
+    CUDA::cudart_static
+)