cuquantum-python-jax-cu12 0.0.5__tar.gz → 0.0.6__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 (48) hide show
  1. cuquantum_python_jax_cu12-0.0.6/MANIFEST.in +44 -0
  2. cuquantum_python_jax_cu12-0.0.6/NV.LICENSE +121 -0
  3. {cuquantum_python_jax_cu12-0.0.5/cuquantum_python_jax_cu12.egg-info → cuquantum_python_jax_cu12-0.0.6}/PKG-INFO +19 -19
  4. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/README.md +14 -15
  5. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/configure.sh +2 -2
  6. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/__init__.py +10 -1
  7. cuquantum_python_jax_cu12-0.0.6/cuquantum/densitymat/jax/_build_info.py +19 -0
  8. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/cppsrc/cudensitymat.h +993 -97
  9. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/cppsrc/cudensitymat_jax.cpp +15 -19
  10. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/cppsrc/cudensitymat_jax.h +1 -8
  11. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/cppsrc/pybind.cpp +1 -7
  12. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/cppsrc/utils.h +1 -1
  13. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/operator_action.py +31 -54
  14. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/pysrc/context.py +39 -18
  15. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/pysrc/elementary_operator.py +24 -10
  16. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/pysrc/matrix_operator.py +12 -6
  17. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/pysrc/operator.py +36 -24
  18. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/pysrc/operator_action_prim.py +80 -63
  19. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/pysrc/operator_term.py +46 -32
  20. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/utils.py +75 -72
  21. cuquantum_python_jax_cu12-0.0.6/cuquantum/stabilizer/jax/__init__.py +35 -0
  22. cuquantum_python_jax_cu12-0.0.6/cuquantum/stabilizer/jax/_build_info.py +19 -0
  23. cuquantum_python_jax_cu12-0.0.6/cuquantum/stabilizer/jax/cppsrc/CMakeLists.txt +73 -0
  24. cuquantum_python_jax_cu12-0.0.6/cuquantum/stabilizer/jax/cppsrc/custabilizer.h +605 -0
  25. cuquantum_python_jax_cu12-0.0.6/cuquantum/stabilizer/jax/cppsrc/custabilizer_jax.cpp +96 -0
  26. cuquantum_python_jax_cu12-0.0.6/cuquantum/stabilizer/jax/cppsrc/custabilizer_jax.h +13 -0
  27. cuquantum_python_jax_cu12-0.0.6/cuquantum/stabilizer/jax/cppsrc/pybind.cpp +27 -0
  28. cuquantum_python_jax_cu12-0.0.6/cuquantum/stabilizer/jax/pysrc/__init__.py +3 -0
  29. cuquantum_python_jax_cu12-0.0.6/cuquantum/stabilizer/jax/pysrc/_backend.py +36 -0
  30. cuquantum_python_jax_cu12-0.0.6/cuquantum/stabilizer/jax/pysrc/_ffi.py +55 -0
  31. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6/cuquantum_python_jax_cu12.egg-info}/PKG-INFO +19 -19
  32. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum_python_jax_cu12.egg-info/SOURCES.txt +12 -0
  33. cuquantum_python_jax_cu12-0.0.6/cuquantum_python_jax_cu12.egg-info/requires.txt +2 -0
  34. cuquantum_python_jax_cu12-0.0.6/cuquantum_python_jax_cu12.egg-info/top_level.txt +3 -0
  35. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/pyproject.toml +7 -7
  36. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/pyproject.toml.template +5 -5
  37. cuquantum_python_jax_cu12-0.0.6/setup.py +300 -0
  38. cuquantum_python_jax_cu12-0.0.5/MANIFEST.in +0 -9
  39. cuquantum_python_jax_cu12-0.0.5/cuquantum_python_jax_cu12.egg-info/requires.txt +0 -3
  40. cuquantum_python_jax_cu12-0.0.5/cuquantum_python_jax_cu12.egg-info/top_level.txt +0 -2
  41. cuquantum_python_jax_cu12-0.0.5/setup.py +0 -117
  42. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/LICENSE +0 -0
  43. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/cppsrc/CMakeLists.txt +0 -0
  44. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/pysrc/__init__.py +0 -0
  45. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum/densitymat/jax/pysrc/base.py +0 -0
  46. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum_python_jax_cu12.egg-info/dependency_links.txt +0 -0
  47. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/cuquantum_python_jax_cu12.egg-info/not-zip-safe +0 -0
  48. {cuquantum_python_jax_cu12-0.0.5 → cuquantum_python_jax_cu12-0.0.6}/setup.cfg +0 -0
@@ -0,0 +1,44 @@
1
+ #
2
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3
+ # SPDX-License-Identifier: BSD-3-Clause
4
+ #
5
+ # Redistribution and use in source and binary forms, with or without
6
+ # modification, are permitted provided that the following conditions are met:
7
+ #
8
+ # 1. Redistributions of source code must retain the above copyright notice, this
9
+ # list of conditions and the following disclaimer.
10
+ #
11
+ # 2. Redistributions in binary form must reproduce the above copyright notice,
12
+ # this list of conditions and the following disclaimer in the documentation
13
+ # and/or other materials provided with the distribution.
14
+ #
15
+ # 3. Neither the name of the copyright holder nor the names of its
16
+ # contributors may be used to endorse or promote products derived from
17
+ # this software without specific prior written permission.
18
+ #
19
+ # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
20
+ # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
21
+ # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
22
+ # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
23
+ # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
24
+ # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
25
+ # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
26
+ # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
27
+ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
28
+ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
29
+ #
30
+
31
+ include pyproject.toml.template
32
+ include configure.sh
33
+ include cuquantum/densitymat/jax/cppsrc/CMakeLists.txt
34
+ include cuquantum/densitymat/jax/cppsrc/cudensitymat.h
35
+ include cuquantum/densitymat/jax/cppsrc/cudensitymat_jax.cpp
36
+ include cuquantum/densitymat/jax/cppsrc/cudensitymat_jax.h
37
+ include cuquantum/densitymat/jax/cppsrc/pybind.cpp
38
+ include cuquantum/densitymat/jax/cppsrc/utils.h
39
+ include cuquantum/stabilizer/jax/cppsrc/CMakeLists.txt
40
+ include cuquantum/stabilizer/jax/cppsrc/custabilizer.h
41
+ include cuquantum/stabilizer/jax/cppsrc/custabilizer_jax.cpp
42
+ include cuquantum/stabilizer/jax/cppsrc/custabilizer_jax.h
43
+ include cuquantum/stabilizer/jax/cppsrc/pybind.cpp
44
+ prune tests*
@@ -0,0 +1,121 @@
1
+ NVIDIA SOFTWARE LICENSE
2
+
3
+ (version August 18, 2021)
4
+
5
+ This license is a legal agreement between you and NVIDIA Corporation ("NVIDIA") and governs your use of the NVIDIA
6
+ cuQuantum software and materials provided hereunder (“SOFTWARE”).
7
+
8
+ This license can be accepted only by an adult of legal age of majority in the country in which the SOFTWARE is used. If you
9
+ are under the legal age of majority, you must ask your parent or legal guardian to consent to this license. By taking delivery
10
+ of the SOFTWARE, you affirm that you have reached the legal age of majority, you accept the terms of this license, and you
11
+ take legal and financial responsibility for the actions of your permitted users.
12
+
13
+ You agree to use the SOFTWARE only for purposes that are permitted by (a) this license, and (b) any applicable law, regulation
14
+ or generally accepted practices or guidelines in the relevant jurisdictions.
15
+
16
+ 1. LICENSE. Subject to the terms of this license, NVIDIA grants you a non-exclusive limited license to: (a) install and use the
17
+ SOFTWARE, and (b) distribute the SOFTWARE subject to the distribution requirements described in this license. NVIDIA
18
+ reserves all rights, title and interest in and to the SOFTWARE not expressly granted to you under this license.
19
+
20
+ 2. DISTRIBUTION REQUIREMENTS. These are the distribution requirements for you to exercise the distribution grant:
21
+ a. The terms under which you distribute the SOFTWARE must be consistent with the terms of this license, including (without
22
+ limitation) terms relating to the license grant and license restrictions and protection of NVIDIA’s intellectual property rights.
23
+ b. You agree to notify NVIDIA in writing of any known or suspected distribution or use of the SOFTWARE not in compliance
24
+ with the requirements of this license, and to enforce the terms of your agreements with respect to distributed SOFTWARE.
25
+
26
+ 3. LIMITATIONS. Your license to use the SOFTWARE is restricted as follows:
27
+ a. You may not reverse engineer, decompile or disassemble, or remove copyright or other proprietary notices from any
28
+ portion of the SOFTWARE or copies of the SOFTWARE.
29
+ b. You may not modify or create derivative works of any portion of the SOFTWARE.
30
+ c. You may not bypass, disable, or circumvent any technical measure, encryption, security, digital rights management or
31
+ authentication mechanism in the SOFTWARE.
32
+ d. You may not use the SOFTWARE in any manner that would cause it to become subject to an open source software license.
33
+ As examples, licenses that require as a condition of use, modification, and/or distribution that the SOFTWARE be (i) disclosed
34
+ or distributed in source code form; (ii) licensed for the purpose of making derivative works; or (iii) redistributable at no charge.
35
+ e. Unless you have an agreement with NVIDIA for this purpose, you may not use the SOFTWARE with any system or
36
+ application where the use or failure of the system or application can reasonably be expected to threaten or result in personal
37
+ injury, death, or catastrophic loss. Examples include use in avionics, navigation, military, medical, life support or other life
38
+ critical applications. NVIDIA does not design, test or manufacture the SOFTWARE for these critical uses and NVIDIA shall not
39
+ be liable to you or any third party, in whole or in part, for any claims or damages arising from such uses.
40
+ f. You agree to defend, indemnify and hold harmless NVIDIA and its affiliates, and their respective employees, contractors,
41
+ agents, officers and directors, from and against any and all claims, damages, obligations, losses, liabilities, costs or debt, fines,
42
+ restitutions and expenses (including but not limited to attorney’s fees and costs incident to establishing the right of
43
+ indemnification) arising out of or related to use of the SOFTWARE outside of the scope of this Agreement, or not in compliance
44
+ with its terms.
45
+
46
+ 4. PRE-RELEASE. SOFTWARE versions identified as alpha, beta, preview, early access or otherwise as pre-release may not be
47
+ fully functional, may contain errors or design flaws, and may have reduced or different security, privacy, availability, and
48
+ reliability standards relative to commercial versions of NVIDIA software and materials. You may use a pre-release SOFTWARE
49
+ version at your own risk, understanding that these versions are not intended for use in production or business-critical
50
+ systems.
51
+
52
+ 5. OWNERSHIP. The SOFTWARE and the related intellectual property rights therein are and will remain the sole and exclusive
53
+ property of NVIDIA or its licensors. The SOFTWARE is copyrighted and protected by the laws of the United States and other
54
+ countries, and international treaty provisions. NVIDIA may make changes to the SOFTWARE, at any time without notice, but
55
+ is not obligated to support or update the SOFTWARE.
56
+
57
+ 6. COMPONENTS UNDER OTHER LICENSES. The SOFTWARE may include NVIDIA or third-party components with separate legal
58
+ notices or terms as may be described in proprietary notices accompanying the SOFTWARE. If and to the extent there is a
59
+ conflict between the terms in this license and the license terms associated with a component, the license terms associated
60
+ with the components control only to the extent necessary to resolve the conflict.
61
+
62
+ 7. FEEDBACK. You may, but don’t have to, provide to NVIDIA any Feedback. “Feedback” means any suggestions, bug fixes,
63
+ enhancements, modifications, feature requests or other feedback regarding the SOFTWARE. For any Feedback that you
64
+ voluntarily provide, you hereby grant NVIDIA and its affiliates a perpetual, non-exclusive, worldwide, irrevocable license to
65
+ use, reproduce, modify, license, sublicense (through multiple tiers of sublicensees), and distribute (through multiple tiers of
66
+ distributors) the Feedback without the payment of any royalties or fees to you. NVIDIA will use Feedback at its choice.
67
+
68
+ 8. NO WARRANTIES. THE SOFTWARE IS PROVIDED "AS IS" WITHOUT ANY EXPRESS OR IMPLIED WARRANTY OF ANY KIND
69
+ INCLUDING, BUT NOT LIMITED TO, WARRANTIES OF MERCHANTABILITY, NONINFRINGEMENT, OR FITNESS FOR A PARTICULAR
70
+ PURPOSE. NVIDIA DOES NOT WARRANT THAT THE SOFTWARE WILL MEET YOUR REQUIREMENTS OR THAT THE OPERATION
71
+ THEREOF WILL BE UNINTERRUPTED OR ERROR-FREE, OR THAT ALL ERRORS WILL BE CORRECTED.
72
+
73
+ 9. LIMITATIONSOF LIABILITY. TO THE MAXIMUM EXTENT PERMITTED BY LAW, NVIDIA AND ITS AFFILIATES SHALL NOT BE LIABLE
74
+ FOR ANY SPECIAL, INCIDENTAL, PUNITIVE OR CONSEQUENTIAL DAMAGES, OR ANY LOST PROFITS, PROJECT DELAYS, LOSS OF
75
+ USE, LOSS OF DATA OR LOSS OF GOODWILL, OR THE COSTS OF PROCURING SUBSTITUTE PRODUCTS, ARISING OUT OF OR IN
76
+ CONNECTION WITH THIS LICENSE OR THE USE OR PERFORMANCE OF THE SOFTWARE, WHETHER SUCH LIABILITY ARISES FROM
77
+ ANY CLAIM BASED UPON BREACH OF CONTRACT, BREACH OF WARRANTY, TORT (INCLUDING NEGLIGENCE), PRODUCT
78
+ LIABILITY OR ANY OTHER CAUSE OF ACTION OR THEORY OF LIABILITY, EVEN IF NVIDIA HAS PREVIOUSLY BEEN ADVISED OF,
79
+ OR COULD REASONABLY HAVE FORESEEN, THE POSSIBILITY OF SUCH DAMAGES. IN NO EVENT WILL NVIDIA’S AND ITS
80
+ AFFILIATES TOTAL CUMULATIVE LIABILITY UNDER OR ARISING OUT OF THIS LICENSE EXCEED US$10.00. THE NATURE OF THE
81
+ LIABILITY OR THE NUMBER OF CLAIMS OR SUITS SHALL NOT ENLARGE OR EXTEND THIS LIMIT.
82
+
83
+ 10. TERMINATION. Your rights under this license will terminate automatically without notice from NVIDIA if you fail to comply
84
+ with any term and condition of this license or if you commence or participate in any legal proceeding against NVIDIA with
85
+ respect to the SOFTWARE. NVIDIA may terminate this license with advance written notice to you if NVIDIA decides to no
86
+ longer provide the SOFTWARE in a country or, in NVIDIA’s sole discretion, the continued use of it is no longer commercially
87
+ viable. Upon any termination of this license, you agree to promptly discontinue use of the SOFTWARE and destroy all copies
88
+ in your possession or control. Your prior distributions in accordance with this license are not affected by the termination of
89
+ this license. All provisions of this license will survive termination, except for the license granted to you.
90
+
91
+ 11. APPLICABLE LAW. This license will be governed in all respects by the laws of the United States and of the State of Delaware
92
+ as those laws are applied to contracts entered into and performed entirely within Delaware by Delaware residents, without
93
+ regard to the conflicts of laws principles. The United Nations Convention on Contracts for the International Sale of Goods is
94
+ specifically disclaimed. You agree to all terms of this Agreement in the English language. The state or federal courts residing
95
+ in Santa Clara County, California shall have exclusive jurisdiction over any dispute or claim arising out of this license.
96
+ Notwithstanding this, you agree that NVIDIA shall still be allowed to apply for injunctive remedies or an equivalent type of
97
+ urgent legal relief in any jurisdiction.
98
+
99
+ 12. NO ASSIGNMENT. This license and your rights and obligations thereunder may not be assigned by you by any means or
100
+ operation of law without NVIDIA’s permission. Any attempted assignment not approved by NVIDIA in writing shall be void
101
+ and of no effect.
102
+
103
+ 13. EXPORT. The SOFTWARE is subject to United States export laws and regulations. You agree that you will not ship, transfer
104
+ or export the SOFTWARE into any country, or use the SOFTWARE in any manner, prohibited by the United States Bureau of
105
+ Industry and Security or economic sanctions regulations administered by the U.S. Department of Treasury’s Office of Foreign
106
+ Assets Control (OFAC), or any applicable export laws, restrictions or regulations. These laws include restrictions on
107
+ destinations, end users and end use. By accepting this license, you confirm that you are not a resident or citizen of any country
108
+ currently embargoed by the U.S. and that you are not otherwise prohibited from receiving the SOFTWARE.
109
+
110
+ 14. GOVERNMENT USE. The SOFTWARE has been developed entirely at private expense and is “commercial items” consisting
111
+ of “commercial computer software” and “commercial computer software documentation” provided with RESTRICTED
112
+ RIGHTS. Use, duplication or disclosure by the U.S. Government or a U.S. Government subcontractor is subject to the
113
+ restrictions in this license pursuant to DFARS 227.7202-3(a) or as set forth in subparagraphs (b)(1) and (2) of the Commercial
114
+ Computer Software - Restricted Rights clause at FAR 52.227-19, as applicable. Contractor/manufacturer is NVIDIA, 2788 San
115
+ Tomas Expressway, Santa Clara, CA 95051.
116
+
117
+ 15. ENTIRE AGREEMENT. This license is the final, complete and exclusive agreement between the parties relating to the
118
+ subject matter of this license and supersedes all prior or contemporaneous understandings and agreements relating to this
119
+ subject matter, whether oral or written. If any court of competent jurisdiction determines that any provision of this license is
120
+ illegal, invalid or unenforceable, the remaining provisions will remain in full force and effect. This license may only be modified
121
+ in a writing signed by an authorized representative of each party.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: cuquantum-python-jax-cu12
3
- Version: 0.0.5
3
+ Version: 0.0.6
4
4
  Summary: NVIDIA cuQuantum Python JAX
5
5
  Author-email: NVIDIA Corporation <cuquantum-python@nvidia.com>
6
6
  License-Expression: BSD-3-Clause
@@ -13,20 +13,21 @@ Classifier: Programming Language :: Python :: 3 :: Only
13
13
  Classifier: Programming Language :: Python :: 3.11
14
14
  Classifier: Programming Language :: Python :: 3.12
15
15
  Classifier: Programming Language :: Python :: 3.13
16
+ Classifier: Programming Language :: Python :: 3.14
16
17
  Classifier: Programming Language :: Python :: Implementation :: CPython
17
18
  Classifier: Environment :: GPU :: NVIDIA CUDA
18
19
  Classifier: Environment :: GPU :: NVIDIA CUDA :: 12
19
20
  Requires-Python: >=3.11.0
20
21
  Description-Content-Type: text/markdown
21
22
  License-File: LICENSE
22
- Requires-Dist: pybind11
23
- Requires-Dist: cuquantum-python-cu12~=26.1.0
24
- Requires-Dist: jax[cuda12-local]<0.7,>=0.5
23
+ License-File: NV.LICENSE
24
+ Requires-Dist: cuquantum-python-cu12~=26.6.0
25
+ Requires-Dist: jax[cuda12-local]>=0.8
25
26
  Dynamic: license-file
26
27
 
27
28
  # cuQuantum Python JAX
28
29
 
29
- 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.
30
+ 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 and the GF(2) sparse-dense matrix multiply from the cuStabilizer library.
30
31
 
31
32
  ## Documentation
32
33
 
@@ -36,17 +37,17 @@ Please visit the [NVIDIA cuQuantum Python documentation](https://docs.nvidia.com
36
37
 
37
38
  ### Requirements
38
39
 
39
- The install-time dependencies of the cuQuantum Python JAX package include:
40
+ The build-time dependencies of the cuQuantum Python JAX package include:
40
41
 
41
- * cuquantum-python-cu12~=26.1.0 for CUDA 12 or cuquantum-python-cu13~=26.1.0 for CUDA 13
42
+ * CUDA Toolkit 12.x or 13.x
42
43
  * jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8,<0.9 for CUDA 13
43
44
  * pybind11
45
+ * wheel
44
46
  * setuptools>=77.0.3
45
47
 
46
48
  Note:
47
- 1. cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13.
48
- 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.
49
- 3. cuQuantum Python JAX wheels are CUDA-versioned: `cuquantum-python-jax-cu12` for CUDA 12 and `cuquantum-python-jax-cu13` for CUDA 13.
49
+ - cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13.
50
+ - cuQuantum Python JAX wheels are CUDA-versioned: `cuquantum-python-jax-cu12` for CUDA 12 and `cuquantum-python-jax-cu13` for CUDA 13.
50
51
 
51
52
  #### Installation using `jax[cudaXX-local]`
52
53
 
@@ -68,26 +69,25 @@ pip install nvidia-cuda-cupti-cu12
68
69
  After installing cuDNN and cuPTI, the user may install cuQuantum Python JAX with `pip` using either:
69
70
 
70
71
  ```
71
- pip install --no-build-isolation cuquantum-python-jax-cu12 # for CUDA 12
72
- pip install --no-build-isolation cuquantum-python-jax-cu13 # for CUDA 13
72
+ pip install cuquantum-python-jax-cu12 # for CUDA 12
73
+ pip install cuquantum-python-jax-cu13 # for CUDA 13
73
74
  ```
74
75
 
75
76
  or one of
76
77
 
77
78
  ```
78
- pip install --no-build-isolation cuquantum-python-cu12[jax]
79
- pip install --no-build-isolation cuquantum-python-cu13[jax]
79
+ pip install cuquantum-python-cu12[jax]
80
+ pip install cuquantum-python-cu13[jax]
80
81
  ```
81
82
 
82
83
  where the CUDA version is explicitly specified on cuquantum-python.
83
84
 
84
85
  Note:
85
- 1. If cuDNN and cuPTI are installed with `pip`, the user does not need to specify library folders in `LD_LIBRARY_PATH`.
86
- 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.
86
+ - If cuDNN and cuPTI are installed with `pip`, the user does not need to specify library folders in `LD_LIBRARY_PATH`.
87
87
 
88
88
  #### Installing from source
89
89
 
90
- 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:
90
+ 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`, run `./configure.sh` to generate a CUDA version-specific `pyproject.toml` from the template, and then:
91
91
 
92
92
  ```
93
93
  pip install .
@@ -102,9 +102,9 @@ The CUDA version is detected automatically from `$CUDA_PATH` and the wheel will
102
102
  Runtime dependencies of the cuQuantum Python JAX package include:
103
103
 
104
104
  * An NVIDIA GPU with compute capability 7.5+
105
- * cuquantum-python-cu12~=26.1.0 for CUDA 12 or cuquantum-python-cu13~=26.1.0 for CUDA 13
105
+ * CUDA Toolkit 12.x or 13.x
106
+ * cuquantum-python-cu12~=26.3.0 for CUDA 12 or cuquantum-python-cu13~=26.3.0 for CUDA 13
106
107
  * jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8,<0.9 for CUDA 13
107
- * pybind11
108
108
 
109
109
  ## Developer Notes
110
110
 
@@ -1,6 +1,6 @@
1
1
  # cuQuantum Python JAX
2
2
 
3
- 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.
3
+ 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 and the GF(2) sparse-dense matrix multiply from the cuStabilizer library.
4
4
 
5
5
  ## Documentation
6
6
 
@@ -10,17 +10,17 @@ Please visit the [NVIDIA cuQuantum Python documentation](https://docs.nvidia.com
10
10
 
11
11
  ### Requirements
12
12
 
13
- The install-time dependencies of the cuQuantum Python JAX package include:
13
+ The build-time dependencies of the cuQuantum Python JAX package include:
14
14
 
15
- * cuquantum-python-cu12~=26.1.0 for CUDA 12 or cuquantum-python-cu13~=26.1.0 for CUDA 13
15
+ * CUDA Toolkit 12.x or 13.x
16
16
  * jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8,<0.9 for CUDA 13
17
17
  * pybind11
18
+ * wheel
18
19
  * setuptools>=77.0.3
19
20
 
20
21
  Note:
21
- 1. cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13.
22
- 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.
23
- 3. cuQuantum Python JAX wheels are CUDA-versioned: `cuquantum-python-jax-cu12` for CUDA 12 and `cuquantum-python-jax-cu13` for CUDA 13.
22
+ - cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13.
23
+ - cuQuantum Python JAX wheels are CUDA-versioned: `cuquantum-python-jax-cu12` for CUDA 12 and `cuquantum-python-jax-cu13` for CUDA 13.
24
24
 
25
25
  #### Installation using `jax[cudaXX-local]`
26
26
 
@@ -42,26 +42,25 @@ pip install nvidia-cuda-cupti-cu12
42
42
  After installing cuDNN and cuPTI, the user may install cuQuantum Python JAX with `pip` using either:
43
43
 
44
44
  ```
45
- pip install --no-build-isolation cuquantum-python-jax-cu12 # for CUDA 12
46
- pip install --no-build-isolation cuquantum-python-jax-cu13 # for CUDA 13
45
+ pip install cuquantum-python-jax-cu12 # for CUDA 12
46
+ pip install cuquantum-python-jax-cu13 # for CUDA 13
47
47
  ```
48
48
 
49
49
  or one of
50
50
 
51
51
  ```
52
- pip install --no-build-isolation cuquantum-python-cu12[jax]
53
- pip install --no-build-isolation cuquantum-python-cu13[jax]
52
+ pip install cuquantum-python-cu12[jax]
53
+ pip install cuquantum-python-cu13[jax]
54
54
  ```
55
55
 
56
56
  where the CUDA version is explicitly specified on cuquantum-python.
57
57
 
58
58
  Note:
59
- 1. If cuDNN and cuPTI are installed with `pip`, the user does not need to specify library folders in `LD_LIBRARY_PATH`.
60
- 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.
59
+ - If cuDNN and cuPTI are installed with `pip`, the user does not need to specify library folders in `LD_LIBRARY_PATH`.
61
60
 
62
61
  #### Installing from source
63
62
 
64
- 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:
63
+ 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`, run `./configure.sh` to generate a CUDA version-specific `pyproject.toml` from the template, and then:
65
64
 
66
65
  ```
67
66
  pip install .
@@ -76,9 +75,9 @@ The CUDA version is detected automatically from `$CUDA_PATH` and the wheel will
76
75
  Runtime dependencies of the cuQuantum Python JAX package include:
77
76
 
78
77
  * An NVIDIA GPU with compute capability 7.5+
79
- * cuquantum-python-cu12~=26.1.0 for CUDA 12 or cuquantum-python-cu13~=26.1.0 for CUDA 13
78
+ * CUDA Toolkit 12.x or 13.x
79
+ * cuquantum-python-cu12~=26.3.0 for CUDA 12 or cuquantum-python-cu13~=26.3.0 for CUDA 13
80
80
  * jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8,<0.9 for CUDA 13
81
- * pybind11
82
81
 
83
82
  ## Developer Notes
84
83
 
@@ -73,11 +73,11 @@ echo "Detected CUDA_VERSION=${CUDA_VERSION_RAW} (major=${CUDA_MAJOR}) from ${CUD
73
73
  # -------------------------------------------------------------------
74
74
  case "${CUDA_MAJOR}" in
75
75
  12)
76
- JAX_VERSION_SPEC=">=0.5,<0.7"
76
+ JAX_VERSION_SPEC=">=0.8"
77
77
  CUDA_CLASSIFIER="Environment :: GPU :: NVIDIA CUDA :: 12"
78
78
  ;;
79
79
  13)
80
- JAX_VERSION_SPEC=">=0.8,<0.9"
80
+ JAX_VERSION_SPEC=">=0.8"
81
81
  CUDA_CLASSIFIER="Environment :: GPU :: NVIDIA CUDA :: 13"
82
82
  ;;
83
83
  *)
@@ -6,8 +6,17 @@ from cuquantum.bindings._internal import cudensitymat as _cudm
6
6
  _cudm._inspect_function_pointers() # for loading libcudensitymat.so
7
7
 
8
8
  import jax
9
+
10
+ try:
11
+ from ._build_info import check_jax_abi as _check_jax_abi
12
+ except ImportError:
13
+ pass
14
+ else:
15
+ _check_jax_abi()
16
+ del _check_jax_abi
17
+
9
18
  if not jax.config.jax_enable_x64:
10
- raise RuntimeError(f"jax_enable_x64 must be set to True to use cuQuantum Python JAX")
19
+ raise RuntimeError("jax_enable_x64 must be set to True to use cuQuantum Python JAX")
11
20
 
12
21
  from .operator_action import operator_action
13
22
  from .pysrc import (
@@ -0,0 +1,19 @@
1
+ # Auto-generated by setup.py at build time; do not edit.
2
+ JAX_BUILD_VERSION = ""
3
+ PROJECT_NAME = "cuquantum-python-jax-cu12"
4
+
5
+ def check_jax_abi() -> None:
6
+ """Warn whenever the runtime jax version differs from the
7
+ build-time jax version. JAX uses effort-based versioning, so
8
+ any difference -- patch bumps included -- may shift the FFI
9
+ capsule layout these bindings were compiled against."""
10
+ if not JAX_BUILD_VERSION:
11
+ return
12
+ import jax, warnings
13
+ if jax.__version__ != JAX_BUILD_VERSION:
14
+ warnings.warn(
15
+ f"JAX version mismatch: {PROJECT_NAME} was built using "
16
+ f"JAX {JAX_BUILD_VERSION}, current runtime is JAX "
17
+ f"{jax.__version__}. Rebuild against the runtime jax if "
18
+ f"you see issues.",
19
+ RuntimeWarning, stacklevel=3)