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.
Files changed (36) hide show
  1. {cuquantum_python_jax-0.0.4 → cuquantum_python_jax-0.0.7}/LICENSE +1 -1
  2. cuquantum_python_jax-0.0.7/PKG-INFO +79 -0
  3. cuquantum_python_jax-0.0.7/README.rst +37 -0
  4. cuquantum_python_jax-0.0.7/cuda_autodetect.py +268 -0
  5. cuquantum_python_jax-0.0.7/cuquantum_python_jax.egg-info/PKG-INFO +79 -0
  6. cuquantum_python_jax-0.0.7/cuquantum_python_jax.egg-info/SOURCES.txt +9 -0
  7. cuquantum_python_jax-0.0.7/cuquantum_python_jax.egg-info/top_level.txt +1 -0
  8. cuquantum_python_jax-0.0.7/setup.py +99 -0
  9. cuquantum_python_jax-0.0.4/MANIFEST.in +0 -6
  10. cuquantum_python_jax-0.0.4/PKG-INFO +0 -123
  11. cuquantum_python_jax-0.0.4/README.md +0 -86
  12. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/__init__.py +0 -18
  13. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/cppsrc/CMakeLists.txt +0 -99
  14. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/cppsrc/cudensitymat_jax.cpp +0 -435
  15. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/cppsrc/cudensitymat_jax.h +0 -13
  16. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/cppsrc/pybind.cpp +0 -35
  17. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/cppsrc/utils.h +0 -31
  18. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/operator_action.py +0 -323
  19. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/__init__.py +0 -8
  20. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/base.py +0 -98
  21. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/context.py +0 -241
  22. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/elementary_operator.py +0 -248
  23. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/matrix_operator.py +0 -188
  24. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/operator.py +0 -343
  25. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/operator_action_prim.py +0 -429
  26. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/pysrc/operator_term.py +0 -509
  27. cuquantum_python_jax-0.0.4/cuquantum/densitymat/jax/utils.py +0 -73
  28. cuquantum_python_jax-0.0.4/cuquantum_python_jax.egg-info/PKG-INFO +0 -123
  29. cuquantum_python_jax-0.0.4/cuquantum_python_jax.egg-info/SOURCES.txt +0 -27
  30. cuquantum_python_jax-0.0.4/cuquantum_python_jax.egg-info/requires.txt +0 -3
  31. cuquantum_python_jax-0.0.4/cuquantum_python_jax.egg-info/top_level.txt +0 -2
  32. cuquantum_python_jax-0.0.4/pyproject.toml +0 -3
  33. cuquantum_python_jax-0.0.4/setup.py +0 -178
  34. {cuquantum_python_jax-0.0.4 → cuquantum_python_jax-0.0.7}/cuquantum_python_jax.egg-info/dependency_links.txt +0 -0
  35. {cuquantum_python_jax-0.0.4 → cuquantum_python_jax-0.0.7}/cuquantum_python_jax.egg-info/not-zip-safe +0 -0
  36. {cuquantum_python_jax-0.0.4 → cuquantum_python_jax-0.0.7}/setup.cfg +0 -0
@@ -1,4 +1,4 @@
1
- Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
 
3
3
  BSD-3-Clause
4
4
 
@@ -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,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*