cuquantum-python-jax 0.0.2__tar.gz → 0.0.3__tar.gz

{cuquantum_python_jax-0.0.2/cuquantum_python_jax.egg-info → cuquantum_python_jax-0.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuquantum-python-jax
-Version: 0.0.2
+Version: 0.0.3
 Summary: NVIDIA cuQuantum Python JAX
 Home-page: https://developer.nvidia.com/cuquantum-sdk
 Author: NVIDIA Corporation
@@ -16,13 +16,13 @@ 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
+Classifier: Environment :: GPU :: NVIDIA CUDA :: 13
 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
+Requires-Dist: cuquantum-python-cu13~=25.11
+Requires-Dist: jax[cuda13-local]<0.9,>=0.8
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier
@@ -49,18 +49,18 @@ Please visit the [NVIDIA cuQuantum Python documentation](https://docs.nvidia.com
 
 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
+* cuquantum-python-cu12~=25.11 for CUDA 12 or cuquantum-python-cu13~=25.11 for CUDA 13
+* jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8,<0.9 for CUDA 13
 * pybind11
+* setuptools>=77.0.3
 
-Note: cuQuantum Python JAX is only supported with CUDA 12.
+Note: cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13.
 
-#### Installation using `jax[cuda12-local]`
+#### Installation using `jax[cudaXX-local]`
 
-`cuquantum-python-jax` depends explicitly on `jax[cuda12-local]`. `pip install cuquantum-python-jax` will install `jax[cuda12-local]`.
+`cuquantum-python-jax` depends explicitly on `jax[cudaXX-local]`. `pip install cuquantum-python-jax` will install `jax[cudaXX-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`.
+Using `jax[cudaXX-local]` assumes the user provides both cuDNN and the CUDA Toolkit. cuDNN is not a part of the CUDA Toolkit and requires an additional installation. The user must also specify `LD_LIBRARY_PATH`, including the library folders containing `libcudnn.so` and `libcupti.so`.
 
 `libcupti.so` is provided by the CUDA Toolkit. If the CUDA Toolkit is installed under `/usr/local/cuda`, `libcupti.so` is located under `/usr/local/cuda/extras/CUPTI/lib64` and `LD_LIBRARY_PATH` should contain this path.
 
@@ -83,32 +83,11 @@ or
 
 ```
 pip install cuquantum-python-cu12[jax]
+pip install cuquantum-python-cu13[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:
@@ -127,8 +106,8 @@ Where `CUDENSITYMAT_ROOT` is the path to the libraries parent directory. For exa
 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
+* cuquantum-python-cu12~=25.11 for CUDA 12 or cuquantum-python-cu13~=25.11 for CUDA 13
+* jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8<0.9 for CUDA 13 
 * pybind11
 
 ## Developer Notes

{cuquantum_python_jax-0.0.2 → cuquantum_python_jax-0.0.3}/README.md

@@ -12,18 +12,18 @@ Please visit the [NVIDIA cuQuantum Python documentation](https://docs.nvidia.com
 
 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
+* cuquantum-python-cu12~=25.11 for CUDA 12 or cuquantum-python-cu13~=25.11 for CUDA 13
+* jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8,<0.9 for CUDA 13
 * pybind11
+* setuptools>=77.0.3
 
-Note: cuQuantum Python JAX is only supported with CUDA 12.
+Note: cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13.
 
-#### Installation using `jax[cuda12-local]`
+#### Installation using `jax[cudaXX-local]`
 
-`cuquantum-python-jax` depends explicitly on `jax[cuda12-local]`. `pip install cuquantum-python-jax` will install `jax[cuda12-local]`.
+`cuquantum-python-jax` depends explicitly on `jax[cudaXX-local]`. `pip install cuquantum-python-jax` will install `jax[cudaXX-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`.
+Using `jax[cudaXX-local]` assumes the user provides both cuDNN and the CUDA Toolkit. cuDNN is not a part of the CUDA Toolkit and requires an additional installation. The user must also specify `LD_LIBRARY_PATH`, including the library folders containing `libcudnn.so` and `libcupti.so`.
 
 `libcupti.so` is provided by the CUDA Toolkit. If the CUDA Toolkit is installed under `/usr/local/cuda`, `libcupti.so` is located under `/usr/local/cuda/extras/CUPTI/lib64` and `LD_LIBRARY_PATH` should contain this path.
 
@@ -46,32 +46,11 @@ or
 
 ```
 pip install cuquantum-python-cu12[jax]
+pip install cuquantum-python-cu13[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:
@@ -90,8 +69,8 @@ Where `CUDENSITYMAT_ROOT` is the path to the libraries parent directory. For exa
 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
+* cuquantum-python-cu12~=25.11 for CUDA 12 or cuquantum-python-cu13~=25.11 for CUDA 13
+* jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8<0.9 for CUDA 13 
 * pybind11
 
 ## Developer Notes

{cuquantum_python_jax-0.0.2 → cuquantum_python_jax-0.0.3}/cuquantum/densitymat/jax/__init__.py

@@ -2,6 +2,13 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
+import warnings
+
+import jax
+
+if not jax.config.jax_enable_x64:
+    raise RuntimeError(f"jax_enable_x64 must be set to True to use cuQuantum Python JAX")
+
 from .operator_action import operator_action
 from .pysrc.operators import (
     ElementaryOperator,

{cuquantum_python_jax-0.0.2 → cuquantum_python_jax-0.0.3}/cuquantum/densitymat/jax/operator_action.py

@@ -19,9 +19,6 @@ from .pysrc.operator_action import (
     OperatorActionBackwardDiffPrimitive
 )
 
-# params are required to be float64 when passed to the cuDensityMat library.
-jax.config.update("jax_enable_x64", True)
-
 logger = logging.getLogger("cudensitymat-jax.operator_action")
 
 
@@ -62,7 +59,9 @@ def operator_action(op: Operator,
     
     # Process input arguments.
     if isinstance(state_in_bufs, jax.Array):
-        state_in_bufs = [state_in_bufs]
+        state_in_bufs = (state_in_bufs,)
+    else:
+        state_in_bufs = tuple(state_in_bufs)
 
     # Check all states are of the same shape.
     state_in_shape = state_in_bufs[0].shape
@@ -71,9 +70,9 @@ def operator_action(op: Operator,
             raise ValueError("All input state buffers must have the same shape.")
 
     # Determine the state purity.
-    if state_in_shape == op.dims:
+    if state_in_shape == op.dims:  # state vector
         purity = cudm.StatePurity.PURE
-    elif state_in_shape == (*op.dims, *op.dims):
+    elif state_in_shape == (*op.dims, *op.dims):  # density matrix
         purity = cudm.StatePurity.MIXED
     else:
         raise ValueError("The dimensions of the input state do not match the dimensions of the operator.")
@@ -108,7 +107,7 @@ def operator_action(op: Operator,
 @jax.custom_vjp
 def _operator_action(op: Operator,
                      t: float,
-                     state_in_bufs: Sequence[jax.Array],
+                     state_in_bufs: Tuple[jax.Array, ...],
                      params: jax.Array
                      ) -> List[jax.Array]:
     """
@@ -121,7 +120,7 @@ def _operator_action(op: Operator,
 
 def _operator_action_fwd(op: Operator,
                          t: float,
-                         state_in_bufs: Sequence[jax.Array],
+                         state_in_bufs: Tuple[jax.Array, ...],
                          params: jax.Array
                          ) -> Tuple[List[jax.Array], tuple]:
     """
@@ -135,6 +134,9 @@ def _operator_action_fwd(op: Operator,
 def _operator_action_bwd(res: tuple, state_out_adj_bufs: jax.Array | Sequence[jax.Array]) -> tuple:
     """
     Backward rule for operator_action.
+
+    Args:
+        state_out_adj_bufs: Data buffers of the output state adjoint.
     """
     logger.info(f"Calling _operator_action_bwd")
 
@@ -147,7 +149,9 @@ def _operator_action_bwd(res: tuple, state_out_adj_bufs: jax.Array | Sequence[ja
 
     # Process input argument.
     if isinstance(state_out_adj_bufs, jax.Array):
-        state_out_adj_bufs = [state_out_adj_bufs]
+        state_out_adj_bufs = (state_out_adj_bufs,)
+    else:
+        state_out_adj_bufs = tuple(state_out_adj_bufs)
 
     if len(state_in_bufs) != len(state_out_adj_bufs):
         raise ValueError("state_in_bufs and state_out_adj_bufs must have the same number of components.")

{cuquantum_python_jax-0.0.2 → cuquantum_python_jax-0.0.3}/cuquantum/densitymat/jax/pysrc/operator_action.py

@@ -154,7 +154,7 @@ register_primitive(OperatorActionPrimitive)
 
 def operator_action_prim(op: Operator,
                          t: float,
-                         state_in_bufs: Sequence[jax.Array],
+                         state_in_bufs: Tuple[jax.Array, ...],
                          params: jax.Array
                          ) -> List[jax.Array]:
     """
@@ -308,8 +308,8 @@ register_primitive(OperatorActionBackwardDiffPrimitive)
 
 def operator_action_backward_diff_prim(op: Operator,
                                        t: float,
-                                       state_in_bufs: Sequence[jax.Array],
-                                       state_out_adj_bufs: Sequence[jax.Array],
+                                       state_in_bufs: Tuple[jax.Array, ...],
+                                       state_out_adj_bufs: Tuple[jax.Array, ...],
                                        params: jax.Array
                                        ) -> Tuple[jax.Array, ...]:
     """

{cuquantum_python_jax-0.0.2 → cuquantum_python_jax-0.0.3}/cuquantum/densitymat/jax/pysrc/operators.py

@@ -254,7 +254,7 @@ class OperatorTerm:
         self.modes: List[Tuple[int, ...]] = []
         self.conjs: List[Tuple[bool, ...]] = []
         self.duals: List[Tuple[bool, ...]] = []
-        self.coeffs: List[float] = []
+        self.coeffs: List[float | complex] = []
         self.coeff_callbacks: List[cudm.WrappedScalarCallback | None] = []
         self.coeff_grad_callbacks: List[cudm.WrappedScalarGradientCallback | None] = []
 
@@ -346,7 +346,7 @@ class OperatorTerm:
                modes: Sequence[int] | None = None,
                conjs: Sequence[bool] | None = None,
                duals: Sequence[bool] | None = None,
-               coeff: float = 1.0,
+               coeff: float | complex = 1.0,
                coeff_callback: cudm.WrappedScalarCallback | None = None,
                coeff_grad_callback: cudm.WrappedScalarGradientCallback | None = None
                ) -> None:
@@ -356,6 +356,7 @@ class OperatorTerm:
         Args:
             op_prod: Product of elementary or matrix operators to be appended.
             modes: Modes acted on by the operator product.
+            conjs: Conjugations in the operator product. Only applies to MatrixOperators.
             duals: Dualities of the operator product.
             coeff: Coefficient of the operator product.
             coeff_callback: Forward callback for the coeffient.
@@ -365,54 +366,76 @@ class OperatorTerm:
         self._check_dtype(op_prod)
         self._check_and_append_op_prod_type(op_prod)
 
-        # TODO: Check that modes have to be in dims.
-
+        # Check consistency and append modes, conjs and duals.
         if self._op_prod_types[-1] is ElementaryOperator:
-            # Check consistency of modes.
+            # Modes have to specified for elementary operators.
             if modes is None:
                 raise ValueError("Modes acted on must be specified for elementary operators.")
-            
-            # FIXME: Check dims of each elementary operator
-            if len(modes) != len(op_prod):
-                pass
-                # raise ValueError(f"Number of modes acted on {len(modes)} does not match number of operator products {len(op_prod)}.")
-            
+
+            # Check all modes are in Hilbert space.
+            if not set(modes) <= set(range(len(self.dims))):
+                raise ValueError("Modes acted on must be in the Hilbert space, i.e. between 0 and len(self.dims) - 1")
+
+            # Check length of modes acted on are the same as combined number of modes in the operator product.
+            if len(modes) != sum([elem_op.num_modes for elem_op in op_prod]):
+                raise ValueError(f"Number of modes acted on {len(modes)} does not match combined number of modes in the operator product.")
+
+            # Check mode extents of each elementary operator match corresponding qubit dimensions.
+            modes_index = 0
+            for elem_op in op_prod:
+                if elem_op.mode_extents != tuple(
+                    [self.dims[modes[i]] for i in range(modes_index, modes_index + elem_op.num_modes)]
+                ):
+                    raise ValueError("Mode extents of each elementary operator must match corresponding qubit dimensions.")
+                modes_index += elem_op.num_modes
+
             # Check that matrix conjugations cannot be specified for elementary operators.
             if conjs is not None:
                 raise ValueError("Matrix conjugations cannot be specified for elementary operators.")
 
+            # Check that number of duals matches number of modes.
+            if duals is None:
+                duals = (False,) * len(modes)
+            else:
+                if len(duals) != len(modes):
+                    raise ValueError("Number of duals must match number of modes acted on for elementary operator product.")
+
+            # For elementary operator product, we only need modes and duals.
+            self.modes.append(tuple(modes))
+            self.conjs.append(())  # empty tuple is appended here to preserve length
+            self.duals.append(tuple(duals))
+
         else:  # matrix operator product
+            # Check that mode extents match Hilbert space dimensions.
+            for matrix_op in op_prod:
+                if matrix_op.mode_extents != self.dims:
+                    raise ValueError("Mode extents must match Hilbert space dimensions for matrix operators.")
+
+            # Check that modes acted on cannot be specified for matrix operators.
+            if modes is not None:
+                raise ValueError("Modes acted on cannot be specified for matrix operators.")
+
             # Check consistency of conjs.
             if conjs is None:
                 conjs = (False,) * len(op_prod)
             else:
-                # FIXME: Check dims of each elementary operator
                 if len(conjs) != len(op_prod):
-                    pass
-                    # raise ValueError("Number of matrix conjugations must match number of operator products.")
+                    raise ValueError("Number of matrix conjugations must match number of operator products.")
 
-            # Check that modes acted on cannot be specified for matrix operators.
-            if modes is not None:
-                raise ValueError("Modes acted on cannot be specified for matrix operators.")
+            # Check that number of duals matches number of matrix operators.
+            if duals is None:
+                duals = (False,) * len(op_prod)
+            else:
+                if len(duals) != len(op_prod):
+                    raise ValueError("Number of duals must match number of matrix operators.")
 
-        if duals is None:
-            duals = (False,) * len(op_prod)
-        else:
-            # FIXME: Check dims of each elementary operator
-            if len(duals) != len(op_prod):
-                pass
-                # raise ValueError("Number of duals must match number of operator products.")
+            # For matrix operator product, we only need conjs and duals.
+            self.modes.append(tuple(range(len(self.dims))))  # used in reference implementation during testing
+            self.conjs.append(tuple(conjs))
+            self.duals.append(tuple(duals))
 
         # Populate instance attributes.
         self.op_prods.append(tuple(op_prod))
-        if modes is not None:
-            self.modes.append(tuple(modes))
-        else:
-            # Append modes here for testing purposes.
-            self.modes.append(tuple(range(len(self.dims))))
-        if conjs is not None:
-            self.conjs.append(tuple(conjs))
-        self.duals.append(tuple(duals))
         self.coeffs.append(coeff)
         self.coeff_callbacks.append(coeff_callback)
         self.coeff_grad_callbacks.append(coeff_grad_callback)
@@ -498,7 +521,7 @@ class Operator:
 
         self.op_terms: List[OperatorTerm] = []
         self.duals: List[bool] = []
-        self.coeffs: List[float] = []
+        self.coeffs: List[float | complex] = []
         self.coeff_callbacks: List[cudm.WrappedScalarCallback | None] = []
         self.coeff_grad_callbacks: List[cudm.WrappedScalarGradientCallback | None] = []
 
@@ -564,7 +587,7 @@ class Operator:
                op_term: OperatorTerm,
                *,
                dual: bool = False,
-               coeff: float = 1.0,
+               coeff: float | complex = 1.0,
                coeff_callback: cudm.WrappedScalarCallback | None = None,
                coeff_grad_callback: cudm.WrappedScalarGradientCallback | None = None
                ) -> None:

{cuquantum_python_jax-0.0.2 → cuquantum_python_jax-0.0.3/cuquantum_python_jax.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuquantum-python-jax
-Version: 0.0.2
+Version: 0.0.3
 Summary: NVIDIA cuQuantum Python JAX
 Home-page: https://developer.nvidia.com/cuquantum-sdk
 Author: NVIDIA Corporation
@@ -16,13 +16,13 @@ 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
+Classifier: Environment :: GPU :: NVIDIA CUDA :: 13
 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
+Requires-Dist: cuquantum-python-cu13~=25.11
+Requires-Dist: jax[cuda13-local]<0.9,>=0.8
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier
@@ -49,18 +49,18 @@ Please visit the [NVIDIA cuQuantum Python documentation](https://docs.nvidia.com
 
 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
+* cuquantum-python-cu12~=25.11 for CUDA 12 or cuquantum-python-cu13~=25.11 for CUDA 13
+* jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8,<0.9 for CUDA 13
 * pybind11
+* setuptools>=77.0.3
 
-Note: cuQuantum Python JAX is only supported with CUDA 12.
+Note: cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13.
 
-#### Installation using `jax[cuda12-local]`
+#### Installation using `jax[cudaXX-local]`
 
-`cuquantum-python-jax` depends explicitly on `jax[cuda12-local]`. `pip install cuquantum-python-jax` will install `jax[cuda12-local]`.
+`cuquantum-python-jax` depends explicitly on `jax[cudaXX-local]`. `pip install cuquantum-python-jax` will install `jax[cudaXX-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`.
+Using `jax[cudaXX-local]` assumes the user provides both cuDNN and the CUDA Toolkit. cuDNN is not a part of the CUDA Toolkit and requires an additional installation. The user must also specify `LD_LIBRARY_PATH`, including the library folders containing `libcudnn.so` and `libcupti.so`.
 
 `libcupti.so` is provided by the CUDA Toolkit. If the CUDA Toolkit is installed under `/usr/local/cuda`, `libcupti.so` is located under `/usr/local/cuda/extras/CUPTI/lib64` and `LD_LIBRARY_PATH` should contain this path.
 
@@ -83,32 +83,11 @@ or
 
 ```
 pip install cuquantum-python-cu12[jax]
+pip install cuquantum-python-cu13[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:
@@ -127,8 +106,8 @@ Where `CUDENSITYMAT_ROOT` is the path to the libraries parent directory. For exa
 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
+* cuquantum-python-cu12~=25.11 for CUDA 12 or cuquantum-python-cu13~=25.11 for CUDA 13
+* jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8<0.9 for CUDA 13 
 * pybind11
 
 ## Developer Notes

cuquantum_python_jax-0.0.3/cuquantum_python_jax.egg-info/requires.txt

@@ -0,0 +1,3 @@
+pybind11
+cuquantum-python-cu13~=25.11
+jax[cuda13-local]<0.9,>=0.8

{cuquantum_python_jax-0.0.2 → cuquantum_python_jax-0.0.3}/setup.py

@@ -7,6 +7,8 @@
 # versioning scheme.
 
 import os
+import re
+import sys
 import shutil
 import pathlib
 import subprocess
@@ -16,6 +18,45 @@ from setuptools.command.build_ext import build_ext
 from setuptools.command.bdist_wheel import bdist_wheel
 
 
+try:
+    cuda_path = os.environ['CUDA_PATH']
+except KeyError as e:
+    raise RuntimeError('CUDA is not found, please set $CUDA_PATH') from e
+
+
+def check_cuda_version():
+    try:
+        # We cannot do a dlopen and call cudaRuntimeGetVersion, because it
+        # requires GPUs. We also do not want to rely on the compiler utility
+        # provided in distutils (deprecated) or setuptools, as this is a very
+        # simple string parsing task.
+        # TODO: switch to cudaRuntimeGetVersion once it's fixed (nvbugs 3624208)
+        cuda_h = os.path.join(cuda_path, 'include', 'cuda.h')
+        with open(cuda_h, 'r') as f:
+            cuda_h = f.read()
+        m = re.search('#define CUDA_VERSION ([0-9]*)', cuda_h)
+        if m:
+            ver = int(m.group(1))
+        else:
+            raise RuntimeError("cannot parse CUDA_VERSION")
+    except:
+        raise
+    else:
+        # 12020 -> "12.2"
+        return str(ver // 1000) + '.' + str((ver % 100) // 10)
+
+
+# We support CUDA 12/13 starting 25.09
+cuda_ver = check_cuda_version()
+
+if '12.0' <= cuda_ver < '13.0':
+    cuda_major_ver = '12'
+elif '13.0' <= cuda_ver < '14.0':
+    cuda_major_ver = '13'
+else:
+    raise RuntimeError(f"Unsupported CUDA version: {cuda_ver}")
+
+
 class CMakeExtension(Extension):
     def __init__(self, name: str, source_dir: str = ""):
         super().__init__(name, sources=[])
@@ -75,7 +116,30 @@ description = "NVIDIA cuQuantum Python JAX"
 with open(os.path.join(os.path.dirname(__file__), "README.md"), encoding="utf-8") as f:
     long_description = f.read()
 
-__version__ = "0.0.2"
+__version__ = "0.0.3"
+
+classifiers=[
+    "Development Status :: 5 - Production/Stable",
+    "Operating System :: POSIX :: Linux",
+    "Topic :: Education",
+    "Topic :: Scientific/Engineering",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: Implementation :: CPython",
+    "Environment :: GPU :: NVIDIA CUDA"
+]
+
+install_requires = ['pybind11']
+if cuda_major_ver == '12':
+    classifiers.append("Environment :: GPU :: NVIDIA CUDA :: 12")
+    install_requires.append("cuquantum-python-cu12~=25.11")
+    install_requires.append("jax[cuda12-local]>=0.5,<0.7")
+else:
+    classifiers.append("Environment :: GPU :: NVIDIA CUDA :: 13")
+    install_requires.append("cuquantum-python-cu13~=25.11")
+    install_requires.append("jax[cuda13-local]>=0.8,<0.9")
 
 setup(
     name="cuquantum-python-jax",
@@ -84,29 +148,13 @@ setup(
     author="NVIDIA Corporation",
     author_email="cuquantum-python@nvidia.com",
     python_requires=">=3.11.0",
-    install_requires=[
-        "cuquantum-python-cu12~=25.09",
-        "jax[cuda12-local]>=0.5,<0.7",
-        "pybind11"
-    ],
+    install_requires=install_requires,
     license="BSD-3-Clause",
     license_files = ('LICENSE',),
     description=description,
     long_description=long_description,
     long_description_content_type="text/markdown",
-    classifiers=[
-        "Development Status :: 5 - Production/Stable",
-        "Operating System :: POSIX :: Linux",
-        "Topic :: Education",
-        "Topic :: Scientific/Engineering",
-        "Programming Language :: Python :: 3 :: Only",
-        "Programming Language :: Python :: 3.11",
-        "Programming Language :: Python :: 3.12",
-        "Programming Language :: Python :: 3.13",
-        "Programming Language :: Python :: Implementation :: CPython",
-        "Environment :: GPU :: NVIDIA CUDA",
-        "Environment :: GPU :: NVIDIA CUDA :: 12",
-    ],
+    classifiers=classifiers,
     ext_modules=[
         CMakeExtension("cppsrc", source_dir="cuquantum/densitymat/jax/cppsrc"),
     ],

cuquantum_python_jax-0.0.2/cuquantum_python_jax.egg-info/requires.txt

@@ -1,3 +0,0 @@
-cuquantum-python-cu12~=25.09
-jax[cuda12-local]<0.7,>=0.5
-pybind11