cuquantum-python-jax 0.0.1__tar.gz

cuquantum_python_jax-0.0.1/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-0.0.1/MANIFEST.in

@@ -0,0 +1,6 @@
+include cuquantum/densitymat/jax/cppsrc/CMakeLists.txt
+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-0.0.1/PKG-INFO

@@ -0,0 +1,138 @@
+Metadata-Version: 2.4
+Name: cuquantum-python-jax
+Version: 0.0.1
+Summary: NVIDIA cuQuantum Python JAX
+Home-page: https://developer.nvidia.com/cuquantum-sdk
+Author: NVIDIA Corporation
+Author-email: cuquantum-python@nvidia.com
+License: BSD-3-Clause
+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: cuquantum-python-cu12~=25.09
+Requires-Dist: jax[cuda12-local]<0.7,>=0.5
+Requires-Dist: pybind11
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# 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 package include:
+
+* cuquantum-python-cu12~=25.09
+* setuptools>=77.0.3
+* jax[cuda12]>=0.5,<0.7 or jax[cuda12-local]>=0.5,<0.7
+* pybind11
+
+Note: cuQuantum Python JAX is only supported with CUDA 12.
+
+#### Installation using `jax[cuda12-local]`
+
+`cuquantum-python-jax` depends explicitly on `jax[cuda12-local]`. `pip install cuquantum-python-jax` will install `jax[cuda12-local]`.
+
+Using `jax[cuda12-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` using `pip` using either:
+
+```
+pip install cuquantum-python-jax
+```
+
+or
+
+```
+pip install cuquantum-python-cu12[jax]
+```
+
+Note: if cuDNN and cuPTI are installed with `pip`, the user does not need to specify library folders in `LD_LIBRARY_PATH`.
+
+#### Installation using `jax[cuda12]`
+
+Alternatively, the user may 
+
+```
+pip install jax[cuda12]  # install cuPTI and cuDNN together with CUDA-enabled JAX
+```
+
+and either 
+
+```
+pip install cuquantum-python-jax
+```
+
+or
+
+```
+pip install cuquantum-python-cu12[jax]
+```
+
+Warning: if the user has an installation of CUDA outside of `pip`, this may create conflicts and undefined behavior.
+
+#### 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:
+
+```
+export CUDENSITYMAT_ROOT=...
+pip install .
+```
+
+Where `CUDENSITYMAT_ROOT` is the path to the libraries parent directory. For example, if `CUDENSITYMAT_ROOT=/usr/local`, `libcudensitymat.so` would be found under `/usr/local/lib` or `/usr/local/lib64`.
+
+## Running
+
+### Requirements
+
+Runtime dependencies of the cuQuantum Python package include:
+
+* An NVIDIA GPU with compute capability 7.5+
+* cuquantum-python-cu12~=25.09
+* jax[cuda12]>=0.5,<0.7 or jax[cuda12-local]>=0.5,<0.7
+* 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-0.0.1/README.md

@@ -0,0 +1,101 @@
+# 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 package include:
+
+* cuquantum-python-cu12~=25.09
+* setuptools>=77.0.3
+* jax[cuda12]>=0.5,<0.7 or jax[cuda12-local]>=0.5,<0.7
+* pybind11
+
+Note: cuQuantum Python JAX is only supported with CUDA 12.
+
+#### Installation using `jax[cuda12-local]`
+
+`cuquantum-python-jax` depends explicitly on `jax[cuda12-local]`. `pip install cuquantum-python-jax` will install `jax[cuda12-local]`.
+
+Using `jax[cuda12-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` using `pip` using either:
+
+```
+pip install cuquantum-python-jax
+```
+
+or
+
+```
+pip install cuquantum-python-cu12[jax]
+```
+
+Note: if cuDNN and cuPTI are installed with `pip`, the user does not need to specify library folders in `LD_LIBRARY_PATH`.
+
+#### Installation using `jax[cuda12]`
+
+Alternatively, the user may 
+
+```
+pip install jax[cuda12]  # install cuPTI and cuDNN together with CUDA-enabled JAX
+```
+
+and either 
+
+```
+pip install cuquantum-python-jax
+```
+
+or
+
+```
+pip install cuquantum-python-cu12[jax]
+```
+
+Warning: if the user has an installation of CUDA outside of `pip`, this may create conflicts and undefined behavior.
+
+#### 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:
+
+```
+export CUDENSITYMAT_ROOT=...
+pip install .
+```
+
+Where `CUDENSITYMAT_ROOT` is the path to the libraries parent directory. For example, if `CUDENSITYMAT_ROOT=/usr/local`, `libcudensitymat.so` would be found under `/usr/local/lib` or `/usr/local/lib64`.
+
+## Running
+
+### Requirements
+
+Runtime dependencies of the cuQuantum Python package include:
+
+* An NVIDIA GPU with compute capability 7.5+
+* cuquantum-python-cu12~=25.09
+* jax[cuda12]>=0.5,<0.7 or jax[cuda12-local]>=0.5,<0.7
+* 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-0.0.1/cuquantum/densitymat/jax/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from .operator_action import operator_action
+from .pysrc.operators import (
+    ElementaryOperator,
+    MatrixOperator,
+    OperatorTerm,
+    Operator
+)

cuquantum_python_jax-0.0.1/cuquantum/densitymat/jax/cppsrc/CMakeLists.txt

@@ -0,0 +1,93 @@
+# 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 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)
+
+# FIXME: This should be made more robust and moved to setup.py.
+execute_process(
+    COMMAND ${Python3_EXECUTABLE} -c "import site; print(f'{site.getsitepackages()[0]}/cuquantum')"
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    OUTPUT_VARIABLE CUQUANTUM_PYTHON_ROOT
+)
+
+if(DEFINED ENV{CUDENSITYMAT_ROOT})
+    set(CUDENSITYMAT_ROOT $ENV{CUDENSITYMAT_ROOT})
+else()
+    set(CUDENSITYMAT_ROOT ${CUQUANTUM_PYTHON_ROOT})
+endif()
+
+message(STATUS "cuQuantum Python directory: ${CUQUANTUM_PYTHON_ROOT}")
+message(STATUS "cuDensityMat directory: ${CUDENSITYMAT_ROOT}")
+
+pybind11_add_module(
+    ${PROJECT_NAME}
+    cudensitymat_jax.cpp
+    pybind.cpp
+)
+target_include_directories(
+    ${PROJECT_NAME}
+    PUBLIC
+    ${CUDAToolkit_INCLUDE_DIRS}
+    ${XLA_DIR}
+    ${pybind11_INCLUDE_DIR}
+    ${CUDENSITYMAT_ROOT}/include
+)
+
+find_library(
+    CUDENSITYMAT_LIBRARY
+    NAMES libcudensitymat.so
+          libcudensitymat.so.0
+    HINTS ${CUDENSITYMAT_ROOT}/lib64
+          ${CUDENSITYMAT_ROOT}/lib
+)
+
+if(NOT CUDENSITYMAT_LIBRARY)
+    message(FATAL_ERROR "cuDensityMat library not found")
+else()
+    message(STATUS "cuDensityMat library: ${CUDENSITYMAT_LIBRARY}")
+endif()
+
+target_link_libraries(
+    ${PROJECT_NAME}
+    PRIVATE
+    CUDA::cudart_static
+    ${CUDENSITYMAT_LIBRARY}
+)