cuquantum-python-jax 0.0.4__tar.gz → 0.0.7__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {cuquantum_python_jax-0.0.4 → cuquantum_python_jax-0.0.7}/LICENSE +1 -1
- cuquantum_python_jax-0.0.7/PKG-INFO +79 -0
- cuquantum_python_jax-0.0.7/README.rst +37 -0
- cuquantum_python_jax-0.0.7/cuda_autodetect.py +268 -0
- cuquantum_python_jax-0.0.7/cuquantum_python_jax.egg-info/PKG-INFO +79 -0
- cuquantum_python_jax-0.0.7/cuquantum_python_jax.egg-info/SOURCES.txt +9 -0
- cuquantum_python_jax-0.0.7/cuquantum_python_jax.egg-info/top_level.txt +1 -0
- cuquantum_python_jax-0.0.7/setup.py +99 -0
- cuquantum_python_jax-0.0.4/MANIFEST.in +0 -6
- cuquantum_python_jax-0.0.4/PKG-INFO +0 -123
- cuquantum_python_jax-0.0.4/README.md +0 -86
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/__init__.py +0 -18
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/cppsrc/CMakeLists.txt +0 -99
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/cppsrc/cudensitymat_jax.cpp +0 -435
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/cppsrc/cudensitymat_jax.h +0 -13
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/cppsrc/pybind.cpp +0 -35
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/cppsrc/utils.h +0 -31
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/operator_action.py +0 -323
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/__init__.py +0 -8
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/base.py +0 -98
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/context.py +0 -241
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/elementary_operator.py +0 -248
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/matrix_operator.py +0 -188
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/operator.py +0 -343
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/operator_action_prim.py +0 -429
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/operator_term.py +0 -509
- cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/utils.py +0 -73
- cuquantum_python_jax-0.0.4/cuquantum_python_jax.egg-info/PKG-INFO +0 -123
- cuquantum_python_jax-0.0.4/cuquantum_python_jax.egg-info/SOURCES.txt +0 -27
- cuquantum_python_jax-0.0.4/cuquantum_python_jax.egg-info/requires.txt +0 -3
- cuquantum_python_jax-0.0.4/cuquantum_python_jax.egg-info/top_level.txt +0 -2
- cuquantum_python_jax-0.0.4/pyproject.toml +0 -3
- cuquantum_python_jax-0.0.4/setup.py +0 -178
- {cuquantum_python_jax-0.0.4 → cuquantum_python_jax-0.0.7}/cuquantum_python_jax.egg-info/dependency_links.txt +0 -0
- {cuquantum_python_jax-0.0.4 → cuquantum_python_jax-0.0.7}/cuquantum_python_jax.egg-info/not-zip-safe +0 -0
- {cuquantum_python_jax-0.0.4 → cuquantum_python_jax-0.0.7}/setup.cfg +0 -0
|
@@ -0,0 +1,79 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: cuquantum-python-jax
|
|
3
|
+
Version: 0.0.7
|
|
4
|
+
Summary: NVIDIA cuQuantum Python JAX
|
|
5
|
+
Home-page: https://developer.nvidia.com/cuquantum-sdk
|
|
6
|
+
Author: NVIDIA Corporation
|
|
7
|
+
Author-email: cuquantum-python@nvidia.com
|
|
8
|
+
License: BSD-3-Clause
|
|
9
|
+
Project-URL: Bug Tracker, https://github.com/NVIDIA/cuQuantum/issues
|
|
10
|
+
Project-URL: User Forum, https://github.com/NVIDIA/cuQuantum/discussions
|
|
11
|
+
Project-URL: Documentation, https://docs.nvidia.com/cuda/cuquantum/latest/python/jax.html
|
|
12
|
+
Project-URL: Source Code, https://github.com/NVIDIA/cuQuantum
|
|
13
|
+
Keywords: cuda,nvidia,state vector,tensor network,high-performance computing,quantum computing
|
|
14
|
+
Classifier: Development Status :: 5 - Production/Stable
|
|
15
|
+
Classifier: Operating System :: POSIX :: Linux
|
|
16
|
+
Classifier: Topic :: Education
|
|
17
|
+
Classifier: Topic :: Scientific/Engineering
|
|
18
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
21
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
22
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
23
|
+
Classifier: Programming Language :: Python :: Implementation :: CPython
|
|
24
|
+
Classifier: Environment :: GPU :: NVIDIA CUDA
|
|
25
|
+
Classifier: Environment :: GPU :: NVIDIA CUDA :: 12
|
|
26
|
+
Classifier: Environment :: GPU :: NVIDIA CUDA :: 13
|
|
27
|
+
Requires-Python: >=3.11
|
|
28
|
+
Description-Content-Type: text/x-rst
|
|
29
|
+
License-File: LICENSE
|
|
30
|
+
Dynamic: author
|
|
31
|
+
Dynamic: author-email
|
|
32
|
+
Dynamic: classifier
|
|
33
|
+
Dynamic: description
|
|
34
|
+
Dynamic: description-content-type
|
|
35
|
+
Dynamic: home-page
|
|
36
|
+
Dynamic: keywords
|
|
37
|
+
Dynamic: license
|
|
38
|
+
Dynamic: license-file
|
|
39
|
+
Dynamic: project-url
|
|
40
|
+
Dynamic: requires-python
|
|
41
|
+
Dynamic: summary
|
|
42
|
+
|
|
43
|
+
**********************
|
|
44
|
+
cuQuantum Python JAX
|
|
45
|
+
**********************
|
|
46
|
+
|
|
47
|
+
cuQuantum Python JAX provides a JAX extension for cuQuantum Python. It exposes selected functionalities of
|
|
48
|
+
cuQuantum SDK in a JAX-compatible way that enables JAX frameworks to directly interface with the exposed
|
|
49
|
+
cuQuantum API. In the current release, cuQuantum JAX exposes JAX interfaces to the Operator Action API from
|
|
50
|
+
the cuDensityMat library and the GF(2) sparse-dense matrix multiplication from the cuStabilizer library.
|
|
51
|
+
|
|
52
|
+
Documentation
|
|
53
|
+
=============
|
|
54
|
+
|
|
55
|
+
For detailed guide, please refer to `cuQuantum Python JAX documentation <https://docs.nvidia.com/cuda/cuquantum/latest/python/jax.html>`_.
|
|
56
|
+
|
|
57
|
+
Installation
|
|
58
|
+
============
|
|
59
|
+
|
|
60
|
+
.. code-block:: bash
|
|
61
|
+
|
|
62
|
+
pip install -v --no-cache-dir cuquantum-python-jax
|
|
63
|
+
|
|
64
|
+
.. note::
|
|
65
|
+
|
|
66
|
+
This package is a meta package pointing to ``cuquantum-python-jax-cuXX``,
|
|
67
|
+
where XX is the CUDA major version (currently CUDA 12 & 13 are supported).
|
|
68
|
+
The meta package will attempt to infer and install the correct ``-cuXX`` wheel.
|
|
69
|
+
The auto-detection mechanism is not guaranteed to work in certain environments, and users are encouraged to install the new wheels that
|
|
70
|
+
come *with* the ``-cuXX`` suffix.
|
|
71
|
+
|
|
72
|
+
The argument ``--no-cache-dir`` is required for pip 23.1+. It forces pip to execute the
|
|
73
|
+
auto-detection logic.
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
Citing cuQuantum
|
|
77
|
+
================
|
|
78
|
+
|
|
79
|
+
`H. Bayraktar et al., "cuQuantum SDK: A High-Performance Library for Accelerating Quantum Science," 2023 IEEE International Conference on Quantum Computing and Engineering (QCE), Bellevue, WA, USA, 2023, pp. 1050-1061, doi: 10.1109/QCE57702.2023.00119 <https://doi.org/10.1109/QCE57702.2023.00119>`_
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
**********************
|
|
2
|
+
cuQuantum Python JAX
|
|
3
|
+
**********************
|
|
4
|
+
|
|
5
|
+
cuQuantum Python JAX provides a JAX extension for cuQuantum Python. It exposes selected functionalities of
|
|
6
|
+
cuQuantum SDK in a JAX-compatible way that enables JAX frameworks to directly interface with the exposed
|
|
7
|
+
cuQuantum API. In the current release, cuQuantum JAX exposes JAX interfaces to the Operator Action API from
|
|
8
|
+
the cuDensityMat library and the GF(2) sparse-dense matrix multiplication from the cuStabilizer library.
|
|
9
|
+
|
|
10
|
+
Documentation
|
|
11
|
+
=============
|
|
12
|
+
|
|
13
|
+
For detailed guide, please refer to `cuQuantum Python JAX documentation <https://docs.nvidia.com/cuda/cuquantum/latest/python/jax.html>`_.
|
|
14
|
+
|
|
15
|
+
Installation
|
|
16
|
+
============
|
|
17
|
+
|
|
18
|
+
.. code-block:: bash
|
|
19
|
+
|
|
20
|
+
pip install -v --no-cache-dir cuquantum-python-jax
|
|
21
|
+
|
|
22
|
+
.. note::
|
|
23
|
+
|
|
24
|
+
This package is a meta package pointing to ``cuquantum-python-jax-cuXX``,
|
|
25
|
+
where XX is the CUDA major version (currently CUDA 12 & 13 are supported).
|
|
26
|
+
The meta package will attempt to infer and install the correct ``-cuXX`` wheel.
|
|
27
|
+
The auto-detection mechanism is not guaranteed to work in certain environments, and users are encouraged to install the new wheels that
|
|
28
|
+
come *with* the ``-cuXX`` suffix.
|
|
29
|
+
|
|
30
|
+
The argument ``--no-cache-dir`` is required for pip 23.1+. It forces pip to execute the
|
|
31
|
+
auto-detection logic.
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
Citing cuQuantum
|
|
35
|
+
================
|
|
36
|
+
|
|
37
|
+
`H. Bayraktar et al., "cuQuantum SDK: A High-Performance Library for Accelerating Quantum Science," 2023 IEEE International Conference on Quantum Computing and Engineering (QCE), Bellevue, WA, USA, 2023, pp. 1050-1061, doi: 10.1109/QCE57702.2023.00119 <https://doi.org/10.1109/QCE57702.2023.00119>`_
|
|
@@ -0,0 +1,268 @@
|
|
|
1
|
+
# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: BSD-3-Clause
|
|
4
|
+
#
|
|
5
|
+
# =============================================================================
|
|
6
|
+
#
|
|
7
|
+
# Heavily influenced by https://github.com/cupy/cupy/blob/main/install/universal_pkg/setup.py.
|
|
8
|
+
# See also the discussion and refs in https://github.com/NVIDIA/cuda-python/issues/16.
|
|
9
|
+
# Below is the original copyright notice from cupy-wheel setup.py.
|
|
10
|
+
#
|
|
11
|
+
# Copyright (c) 2015-2023 Preferred Networks, Inc.
|
|
12
|
+
#
|
|
13
|
+
# SPDX-License-Identifier: MIT
|
|
14
|
+
#
|
|
15
|
+
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
16
|
+
# of this software and associated documentation files (the "Software"), to deal
|
|
17
|
+
# in the Software without restriction, including without limitation the rights
|
|
18
|
+
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
19
|
+
# copies of the Software, and to permit persons to whom the Software is
|
|
20
|
+
# furnished to do so, subject to the following conditions:
|
|
21
|
+
#
|
|
22
|
+
# The above copyright notice and this permission notice shall be included in
|
|
23
|
+
# all copies or substantial portions of the Software.
|
|
24
|
+
#
|
|
25
|
+
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
26
|
+
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
27
|
+
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
28
|
+
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
29
|
+
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
30
|
+
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
|
31
|
+
# THE SOFTWARE.
|
|
32
|
+
|
|
33
|
+
import ctypes
|
|
34
|
+
import importlib.metadata
|
|
35
|
+
import os
|
|
36
|
+
import sys
|
|
37
|
+
from typing import Dict, List, Optional
|
|
38
|
+
|
|
39
|
+
try:
|
|
40
|
+
from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
|
|
41
|
+
except ImportError:
|
|
42
|
+
_bdist_wheel = None
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
# This script is a general utility script, the variables below are to be set by
|
|
46
|
+
# the caller
|
|
47
|
+
PACKAGE_NAME = ''
|
|
48
|
+
PACKAGE_SUPPORTED_CUDA_VER = []
|
|
49
|
+
|
|
50
|
+
# ========================================================================
|
|
51
|
+
|
|
52
|
+
PACKAGE_RESOLUTION = None
|
|
53
|
+
CUDA_RESOLUTION = None
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
class AutoDetectionFailed(Exception):
|
|
57
|
+
def __str__(self) -> str:
|
|
58
|
+
return f'''
|
|
59
|
+
\n\n============================================================
|
|
60
|
+
{super().__str__()}
|
|
61
|
+
============================================================\n
|
|
62
|
+
'''
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def _log(msg: str) -> None:
|
|
66
|
+
sys.stdout.write(f'[{PACKAGE_NAME}] {msg}\n')
|
|
67
|
+
sys.stdout.flush()
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def _get_version_from_library(
|
|
71
|
+
libnames: List[str],
|
|
72
|
+
funcname: str,
|
|
73
|
+
nvrtc: bool = False,
|
|
74
|
+
) -> Optional[int]:
|
|
75
|
+
"""Returns the library version from list of candidate libraries."""
|
|
76
|
+
|
|
77
|
+
for libname in libnames:
|
|
78
|
+
try:
|
|
79
|
+
_log(f'Looking for library: {libname}')
|
|
80
|
+
runtime_so = ctypes.CDLL(libname)
|
|
81
|
+
break
|
|
82
|
+
except Exception as e:
|
|
83
|
+
_log(f'Failed to open {libname}: {e}')
|
|
84
|
+
else:
|
|
85
|
+
_log('No more candidate library to find')
|
|
86
|
+
return None
|
|
87
|
+
|
|
88
|
+
func = getattr(runtime_so, funcname, None)
|
|
89
|
+
if func is None:
|
|
90
|
+
raise AutoDetectionFailed(
|
|
91
|
+
f'{libname}: {func} could not be found')
|
|
92
|
+
func.restype = ctypes.c_int
|
|
93
|
+
|
|
94
|
+
if nvrtc:
|
|
95
|
+
# nvrtcVersion
|
|
96
|
+
func.argtypes = [
|
|
97
|
+
ctypes.POINTER(ctypes.c_int),
|
|
98
|
+
ctypes.POINTER(ctypes.c_int),
|
|
99
|
+
]
|
|
100
|
+
major = ctypes.c_int()
|
|
101
|
+
minor = ctypes.c_int()
|
|
102
|
+
retval = func(major, minor)
|
|
103
|
+
version = major.value * 1000 + minor.value * 10
|
|
104
|
+
else:
|
|
105
|
+
# cudaRuntimeGetVersion
|
|
106
|
+
func.argtypes = [
|
|
107
|
+
ctypes.POINTER(ctypes.c_int),
|
|
108
|
+
]
|
|
109
|
+
version_ref = ctypes.c_int()
|
|
110
|
+
retval = func(version_ref)
|
|
111
|
+
version = version_ref.value
|
|
112
|
+
|
|
113
|
+
if retval != 0: # NVRTC_SUCCESS or cudaSuccess
|
|
114
|
+
raise AutoDetectionFailed(
|
|
115
|
+
f'{libname}: {func} returned error: {retval}')
|
|
116
|
+
_log(f'Detected version: {version}')
|
|
117
|
+
return version
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def _get_cuda_version() -> Optional[int]:
|
|
121
|
+
"""Returns the detected CUDA version or None."""
|
|
122
|
+
|
|
123
|
+
version = None
|
|
124
|
+
|
|
125
|
+
# First try NVRTC
|
|
126
|
+
libnames = [
|
|
127
|
+
'libnvrtc.so.13',
|
|
128
|
+
'libnvrtc.so.12',
|
|
129
|
+
'libnvrtc.so.11.2',
|
|
130
|
+
'libnvrtc.so.11.1',
|
|
131
|
+
'libnvrtc.so.11.0',
|
|
132
|
+
]
|
|
133
|
+
_log(f'Trying to detect CUDA version from libraries: {libnames}')
|
|
134
|
+
try:
|
|
135
|
+
version = _get_version_from_library(libnames, 'nvrtcVersion', True)
|
|
136
|
+
except Exception as e:
|
|
137
|
+
_log(f"Error: {e}") # log and move on
|
|
138
|
+
if version is not None:
|
|
139
|
+
return version
|
|
140
|
+
|
|
141
|
+
# Next try CUDART
|
|
142
|
+
libnames = [
|
|
143
|
+
'libcudart.so.13',
|
|
144
|
+
'libcudart.so.12',
|
|
145
|
+
'libcudart.so.11.0', # side-effect: a CUDA context would be initialized
|
|
146
|
+
]
|
|
147
|
+
_log(f'Trying to detect CUDA version from libraries: {libnames}')
|
|
148
|
+
try:
|
|
149
|
+
version = _get_version_from_library(libnames, 'cudaRuntimeGetVersion', False)
|
|
150
|
+
except Exception as e:
|
|
151
|
+
_log(f"Error: {e}") # log and move on
|
|
152
|
+
if version is not None:
|
|
153
|
+
return version
|
|
154
|
+
|
|
155
|
+
_log("Autodetection failed")
|
|
156
|
+
return None
|
|
157
|
+
|
|
158
|
+
|
|
159
|
+
def _find_installed_packages() -> List[str]:
|
|
160
|
+
"""Returns the list of out packages installed in the environment."""
|
|
161
|
+
|
|
162
|
+
f = lambda x: ''.join([f"{PACKAGE_NAME}-cu", x])
|
|
163
|
+
found = []
|
|
164
|
+
|
|
165
|
+
for pkg in list(map(f, PACKAGE_SUPPORTED_CUDA_VER)):
|
|
166
|
+
try:
|
|
167
|
+
importlib.metadata.distribution(pkg)
|
|
168
|
+
found.append(pkg)
|
|
169
|
+
except importlib.metadata.PackageNotFoundError:
|
|
170
|
+
pass
|
|
171
|
+
return found
|
|
172
|
+
|
|
173
|
+
|
|
174
|
+
def _cuda_version_to_package(ver: int) -> str:
|
|
175
|
+
# TODO: Don't hard-code 11/12, use PACKAGE_SUPPORTED_CUDA_VER instead?
|
|
176
|
+
if ver < 11000:
|
|
177
|
+
raise AutoDetectionFailed(
|
|
178
|
+
f'Your CUDA version ({ver}) is too old.')
|
|
179
|
+
elif ver < 12000:
|
|
180
|
+
# CUDA 11.x
|
|
181
|
+
raise RuntimeError(f'CUDA 11.x is not supported. Please install CUDA 12.x or later.')
|
|
182
|
+
elif ver < 13000:
|
|
183
|
+
# CUDA 12.x
|
|
184
|
+
suffix = '12'
|
|
185
|
+
elif ver < 14000:
|
|
186
|
+
# CUDA 13.x
|
|
187
|
+
suffix = '13'
|
|
188
|
+
else:
|
|
189
|
+
raise AutoDetectionFailed(
|
|
190
|
+
f'Your CUDA version ({ver}) is too new.')
|
|
191
|
+
return f'{PACKAGE_NAME}-cu{suffix}'
|
|
192
|
+
|
|
193
|
+
|
|
194
|
+
# ========================================================================
|
|
195
|
+
|
|
196
|
+
# "Public" API to the caller
|
|
197
|
+
def infer_best_package(package_name: str,
|
|
198
|
+
package_supported_cuda_ver: List[str] = ['12', '13']) -> str:
|
|
199
|
+
"""Returns the appropriate wheel name for the environment."""
|
|
200
|
+
|
|
201
|
+
global PACKAGE_NAME, PACKAGE_SUPPORTED_CUDA_VER
|
|
202
|
+
PACKAGE_NAME = package_name
|
|
203
|
+
PACKAGE_SUPPORTED_CUDA_VER = sorted(package_supported_cuda_ver)
|
|
204
|
+
|
|
205
|
+
# Find the existing wheel installation
|
|
206
|
+
installed = _find_installed_packages()
|
|
207
|
+
|
|
208
|
+
# Detect CUDA version
|
|
209
|
+
version = _get_cuda_version()
|
|
210
|
+
if version is not None:
|
|
211
|
+
to_install = _cuda_version_to_package(version)
|
|
212
|
+
else:
|
|
213
|
+
# TODO: change this in the future
|
|
214
|
+
message = (
|
|
215
|
+
"See below for the error message and instruction.\n\n\n" +
|
|
216
|
+
"************************************************************************\n" +
|
|
217
|
+
"ERROR: Unable to detect NVIDIA CUDA Toolkit installation.\n" +
|
|
218
|
+
"ERROR: If CUDA Toolkit is not installed, please install it first.\n" +
|
|
219
|
+
"ERROR: If CUDA Toolkit is installed but not detected, please explicitly specify the version and run\n" +
|
|
220
|
+
f"ERROR: `pip install {PACKAGE_NAME}-cuXX`, with XX being the major\n" +
|
|
221
|
+
"ERROR: version of your CUDA Toolkit installation.\n" +
|
|
222
|
+
"************************************************************************\n\n"
|
|
223
|
+
)
|
|
224
|
+
raise AutoDetectionFailed(message)
|
|
225
|
+
|
|
226
|
+
# Disallow -cu12 & -cu13 wheels from coexisting
|
|
227
|
+
if len(installed) > 1 or (len(installed) == 1 and installed[0] != to_install):
|
|
228
|
+
raise AutoDetectionFailed(
|
|
229
|
+
f'You already have the {PACKAGE_NAME} package(s) installed: \n'
|
|
230
|
+
f' {installed}\n'
|
|
231
|
+
'while you attempt to install:\n'
|
|
232
|
+
f' {to_install}\n'
|
|
233
|
+
'Please uninstall all of them first, then try reinstalling.')
|
|
234
|
+
|
|
235
|
+
global PACKAGE_RESOLUTION, CUDA_RESOLUTION
|
|
236
|
+
PACKAGE_RESOLUTION = to_install
|
|
237
|
+
CUDA_RESOLUTION = version
|
|
238
|
+
_log(f"Installing {to_install}...")
|
|
239
|
+
return to_install
|
|
240
|
+
|
|
241
|
+
|
|
242
|
+
# "Public" API to the caller
|
|
243
|
+
if _bdist_wheel is not None:
|
|
244
|
+
|
|
245
|
+
# Technically we need a way to force reinstalling the sdist and ignored the cached wheel.
|
|
246
|
+
# That said, I cannot reproduce the past-known behavior in my env. The sdist is always
|
|
247
|
+
# reinstalled, though it could be due to certain combination of pip/setuptools/wheel/etc,
|
|
248
|
+
# and it's still better to keep this WAR.
|
|
249
|
+
|
|
250
|
+
class bdist_wheel(_bdist_wheel):
|
|
251
|
+
|
|
252
|
+
# Adopted from https://discuss.python.org/t/wheel-caching-and-non-deterministic-builds/7687
|
|
253
|
+
|
|
254
|
+
def finalize_options(self):
|
|
255
|
+
super().finalize_options()
|
|
256
|
+
|
|
257
|
+
# Use "cuXX" as the build tag to force re-running sdist if the
|
|
258
|
+
# CUDA version in the user env has changed
|
|
259
|
+
if PACKAGE_RESOLUTION is None:
|
|
260
|
+
assert False, "something went wrong"
|
|
261
|
+
build_tag = PACKAGE_RESOLUTION.split("-")[-1]
|
|
262
|
+
|
|
263
|
+
# per PEP 427, build tag must start with a digit
|
|
264
|
+
self.build_number = f"0_{build_tag}"
|
|
265
|
+
|
|
266
|
+
else:
|
|
267
|
+
|
|
268
|
+
bdist_wheel = None
|
|
@@ -0,0 +1,79 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: cuquantum-python-jax
|
|
3
|
+
Version: 0.0.7
|
|
4
|
+
Summary: NVIDIA cuQuantum Python JAX
|
|
5
|
+
Home-page: https://developer.nvidia.com/cuquantum-sdk
|
|
6
|
+
Author: NVIDIA Corporation
|
|
7
|
+
Author-email: cuquantum-python@nvidia.com
|
|
8
|
+
License: BSD-3-Clause
|
|
9
|
+
Project-URL: Bug Tracker, https://github.com/NVIDIA/cuQuantum/issues
|
|
10
|
+
Project-URL: User Forum, https://github.com/NVIDIA/cuQuantum/discussions
|
|
11
|
+
Project-URL: Documentation, https://docs.nvidia.com/cuda/cuquantum/latest/python/jax.html
|
|
12
|
+
Project-URL: Source Code, https://github.com/NVIDIA/cuQuantum
|
|
13
|
+
Keywords: cuda,nvidia,state vector,tensor network,high-performance computing,quantum computing
|
|
14
|
+
Classifier: Development Status :: 5 - Production/Stable
|
|
15
|
+
Classifier: Operating System :: POSIX :: Linux
|
|
16
|
+
Classifier: Topic :: Education
|
|
17
|
+
Classifier: Topic :: Scientific/Engineering
|
|
18
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
21
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
22
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
23
|
+
Classifier: Programming Language :: Python :: Implementation :: CPython
|
|
24
|
+
Classifier: Environment :: GPU :: NVIDIA CUDA
|
|
25
|
+
Classifier: Environment :: GPU :: NVIDIA CUDA :: 12
|
|
26
|
+
Classifier: Environment :: GPU :: NVIDIA CUDA :: 13
|
|
27
|
+
Requires-Python: >=3.11
|
|
28
|
+
Description-Content-Type: text/x-rst
|
|
29
|
+
License-File: LICENSE
|
|
30
|
+
Dynamic: author
|
|
31
|
+
Dynamic: author-email
|
|
32
|
+
Dynamic: classifier
|
|
33
|
+
Dynamic: description
|
|
34
|
+
Dynamic: description-content-type
|
|
35
|
+
Dynamic: home-page
|
|
36
|
+
Dynamic: keywords
|
|
37
|
+
Dynamic: license
|
|
38
|
+
Dynamic: license-file
|
|
39
|
+
Dynamic: project-url
|
|
40
|
+
Dynamic: requires-python
|
|
41
|
+
Dynamic: summary
|
|
42
|
+
|
|
43
|
+
**********************
|
|
44
|
+
cuQuantum Python JAX
|
|
45
|
+
**********************
|
|
46
|
+
|
|
47
|
+
cuQuantum Python JAX provides a JAX extension for cuQuantum Python. It exposes selected functionalities of
|
|
48
|
+
cuQuantum SDK in a JAX-compatible way that enables JAX frameworks to directly interface with the exposed
|
|
49
|
+
cuQuantum API. In the current release, cuQuantum JAX exposes JAX interfaces to the Operator Action API from
|
|
50
|
+
the cuDensityMat library and the GF(2) sparse-dense matrix multiplication from the cuStabilizer library.
|
|
51
|
+
|
|
52
|
+
Documentation
|
|
53
|
+
=============
|
|
54
|
+
|
|
55
|
+
For detailed guide, please refer to `cuQuantum Python JAX documentation <https://docs.nvidia.com/cuda/cuquantum/latest/python/jax.html>`_.
|
|
56
|
+
|
|
57
|
+
Installation
|
|
58
|
+
============
|
|
59
|
+
|
|
60
|
+
.. code-block:: bash
|
|
61
|
+
|
|
62
|
+
pip install -v --no-cache-dir cuquantum-python-jax
|
|
63
|
+
|
|
64
|
+
.. note::
|
|
65
|
+
|
|
66
|
+
This package is a meta package pointing to ``cuquantum-python-jax-cuXX``,
|
|
67
|
+
where XX is the CUDA major version (currently CUDA 12 & 13 are supported).
|
|
68
|
+
The meta package will attempt to infer and install the correct ``-cuXX`` wheel.
|
|
69
|
+
The auto-detection mechanism is not guaranteed to work in certain environments, and users are encouraged to install the new wheels that
|
|
70
|
+
come *with* the ``-cuXX`` suffix.
|
|
71
|
+
|
|
72
|
+
The argument ``--no-cache-dir`` is required for pip 23.1+. It forces pip to execute the
|
|
73
|
+
auto-detection logic.
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
Citing cuQuantum
|
|
77
|
+
================
|
|
78
|
+
|
|
79
|
+
`H. Bayraktar et al., "cuQuantum SDK: A High-Performance Library for Accelerating Quantum Science," 2023 IEEE International Conference on Quantum Computing and Engineering (QCE), Bellevue, WA, USA, 2023, pp. 1050-1061, doi: 10.1109/QCE57702.2023.00119 <https://doi.org/10.1109/QCE57702.2023.00119>`_
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
README.rst
|
|
3
|
+
cuda_autodetect.py
|
|
4
|
+
setup.py
|
|
5
|
+
cuquantum_python_jax.egg-info/PKG-INFO
|
|
6
|
+
cuquantum_python_jax.egg-info/SOURCES.txt
|
|
7
|
+
cuquantum_python_jax.egg-info/dependency_links.txt
|
|
8
|
+
cuquantum_python_jax.egg-info/not-zip-safe
|
|
9
|
+
cuquantum_python_jax.egg-info/top_level.txt
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
cuda_autodetect
|
|
@@ -0,0 +1,99 @@
|
|
|
1
|
+
# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES.
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: BSD-3-Clause
|
|
4
|
+
|
|
5
|
+
import glob
|
|
6
|
+
import os
|
|
7
|
+
import shutil
|
|
8
|
+
import site
|
|
9
|
+
import subprocess
|
|
10
|
+
import sys
|
|
11
|
+
|
|
12
|
+
from setuptools import setup
|
|
13
|
+
|
|
14
|
+
import cuda_autodetect
|
|
15
|
+
from cuda_autodetect import infer_best_package, bdist_wheel
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
# Update this for every release
|
|
19
|
+
# Must pin to the cuquantum-python-jax-cuXX version
|
|
20
|
+
package_ver = "0.0.7"
|
|
21
|
+
package_name = "cuquantum-python-jax"
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
# get project long description
|
|
25
|
+
with open("README.rst") as f:
|
|
26
|
+
long_description = f.read()
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
# This setup.py handles 2 cases:
|
|
30
|
+
# 1. At the release time, we use it to generate sdist (which contains this script)
|
|
31
|
+
# 2. At the install time, users install the sdist from PyPI
|
|
32
|
+
# and the two cases have different requirements. We distinguish them by
|
|
33
|
+
# setting CUQUANTUM_META_WHEEL_BUILD=1 for Case 1.
|
|
34
|
+
if os.environ.get('CUQUANTUM_META_WHEEL_BUILD', '0') == '1':
|
|
35
|
+
# Case 1: generate sdist
|
|
36
|
+
install_requires = []
|
|
37
|
+
data_files = [('', ['cuda_autodetect.py',])] # extra files to be copied into sdist
|
|
38
|
+
cmdclass = {}
|
|
39
|
+
else:
|
|
40
|
+
# Case 2: install sdist
|
|
41
|
+
install_requires = [f"{infer_best_package(package_name)}=={package_ver}",]
|
|
42
|
+
if 12000 <= cuda_autodetect.CUDA_RESOLUTION < 13000:
|
|
43
|
+
# no ambiguity for CUDA 12, cuquantum-python-jax-cu12 would do the right thing
|
|
44
|
+
pass
|
|
45
|
+
elif 13000 <= cuda_autodetect.CUDA_RESOLUTION < 14000:
|
|
46
|
+
# no ambiguity for CUDA 13, cuquantum-python-jax-cu13 would do the right thing
|
|
47
|
+
pass
|
|
48
|
+
else:
|
|
49
|
+
raise ValueError(f"Unsupported CUDA version: {cuda_autodetect.CUDA_RESOLUTION}")
|
|
50
|
+
data_files = []
|
|
51
|
+
cmdclass = {'bdist_wheel': bdist_wheel} if bdist_wheel is not None else {}
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
setup(
|
|
55
|
+
name=package_name,
|
|
56
|
+
version=package_ver,
|
|
57
|
+
description="NVIDIA cuQuantum Python JAX",
|
|
58
|
+
long_description=long_description,
|
|
59
|
+
long_description_content_type="text/x-rst",
|
|
60
|
+
url="https://developer.nvidia.com/cuquantum-sdk",
|
|
61
|
+
project_urls={
|
|
62
|
+
"Bug Tracker": "https://github.com/NVIDIA/cuQuantum/issues",
|
|
63
|
+
"User Forum": "https://github.com/NVIDIA/cuQuantum/discussions",
|
|
64
|
+
"Documentation": "https://docs.nvidia.com/cuda/cuquantum/latest/python/jax.html",
|
|
65
|
+
"Source Code": "https://github.com/NVIDIA/cuQuantum",
|
|
66
|
+
},
|
|
67
|
+
author="NVIDIA Corporation",
|
|
68
|
+
author_email="cuquantum-python@nvidia.com",
|
|
69
|
+
license="BSD-3-Clause",
|
|
70
|
+
license_files = ('LICENSE',),
|
|
71
|
+
keywords=["cuda", "nvidia", "state vector", "tensor network", "high-performance computing", "quantum computing"],
|
|
72
|
+
# Install files indicated by MANIFEST.in
|
|
73
|
+
# See https://github.com/pypa/sampleproject/issues/30#issuecomment-143947944
|
|
74
|
+
include_package_data=True,
|
|
75
|
+
zip_safe=False,
|
|
76
|
+
data_files=data_files,
|
|
77
|
+
setup_requires=[
|
|
78
|
+
"setuptools",
|
|
79
|
+
"wheel",
|
|
80
|
+
],
|
|
81
|
+
install_requires=install_requires,
|
|
82
|
+
python_requires=">=3.11",
|
|
83
|
+
classifiers=[
|
|
84
|
+
"Development Status :: 5 - Production/Stable",
|
|
85
|
+
"Operating System :: POSIX :: Linux",
|
|
86
|
+
"Topic :: Education",
|
|
87
|
+
"Topic :: Scientific/Engineering",
|
|
88
|
+
"Programming Language :: Python :: 3 :: Only",
|
|
89
|
+
"Programming Language :: Python :: 3.11",
|
|
90
|
+
"Programming Language :: Python :: 3.12",
|
|
91
|
+
"Programming Language :: Python :: 3.13",
|
|
92
|
+
"Programming Language :: Python :: 3.14",
|
|
93
|
+
"Programming Language :: Python :: Implementation :: CPython",
|
|
94
|
+
"Environment :: GPU :: NVIDIA CUDA",
|
|
95
|
+
"Environment :: GPU :: NVIDIA CUDA :: 12",
|
|
96
|
+
"Environment :: GPU :: NVIDIA CUDA :: 13",
|
|
97
|
+
],
|
|
98
|
+
cmdclass=cmdclass,
|
|
99
|
+
)
|
|
@@ -1,6 +0,0 @@
|
|
|
1
|
-
include cuquantum/densitymat/jax/cppsrc/CMakeLists.txt
|
|
2
|
-
include cuquantum/densitymat/jax/cppsrc/cudensitymat_jax.cpp
|
|
3
|
-
include cuquantum/densitymat/jax/cppsrc/cudensitymat_jax.h
|
|
4
|
-
include cuquantum/densitymat/jax/cppsrc/pybind.cpp
|
|
5
|
-
include cuquantum/densitymat/jax/cppsrc/utils.h
|
|
6
|
-
prune tests*
|