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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuquantum-python-jax
-Version: 0.0.2
+Version: 0.0.4
 Summary: NVIDIA cuQuantum Python JAX
 Home-page: https://developer.nvidia.com/cuquantum-sdk
 Author: NVIDIA Corporation
@@ -20,9 +20,9 @@ Classifier: Environment :: GPU :: NVIDIA CUDA :: 12
 Requires-Python: >=3.11.0
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: cuquantum-python-cu12~=25.09
-Requires-Dist: jax[cuda12-local]<0.7,>=0.5
 Requires-Dist: pybind11
+Requires-Dist: cuquantum-python-cu12~=26.1.0
+Requires-Dist: jax[cuda12-local]<0.7,>=0.5
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier
@@ -47,20 +47,22 @@ Please visit the [NVIDIA cuQuantum Python documentation](https://docs.nvidia.com
 
 ### Requirements
 
-The install-time dependencies of the cuQuantum Python package include:
+The install-time dependencies of the cuQuantum Python JAX 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~=26.1.0 for CUDA 12 or cuquantum-python-cu13~=26.1.0 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: 
+1. cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13.
+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.
 
-#### 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.
 
@@ -73,41 +75,24 @@ pip install nvidia-cudnn-cu12
 pip install nvidia-cuda-cupti-cu12
 ```
 
-After installing cuDNN and cuPTI, the user may install `cuquantum-python-jax` using `pip` using either:
+After installing cuDNN and cuPTI, the user may install `cuquantum-python-jax` with `pip` using either:
 
 ```
-pip install cuquantum-python-jax
+pip install --no-build-isolation cuquantum-python-jax
 ```
 
-or
+in which case the CUDA version will be detected, or one of
 
 ```
-pip install cuquantum-python-cu12[jax]
+pip install --no-build-isolation cuquantum-python-cu12[jax]
+pip install --no-build-isolation 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]
-```
+where the CUDA version is explicitly specified on cuquantum-python.
 
-Warning: if the user has an installation of CUDA outside of `pip`, this may create conflicts and undefined behavior.
+Note:
+1. If cuDNN and cuPTI are installed with `pip`, the user does not need to specify library folders in `LD_LIBRARY_PATH`.
+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.
 
 #### Installing from source
 
@@ -127,8 +112,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~=26.1.0 for CUDA 12 or cuquantum-python-cu13~=26.1.0 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.4}/README.md

@@ -10,20 +10,22 @@ Please visit the [NVIDIA cuQuantum Python documentation](https://docs.nvidia.com
 
 ### Requirements
 
-The install-time dependencies of the cuQuantum Python package include:
+The install-time dependencies of the cuQuantum Python JAX 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~=26.1.0 for CUDA 12 or cuquantum-python-cu13~=26.1.0 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: 
+1. cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13.
+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.
 
-#### 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.
 
@@ -36,41 +38,24 @@ pip install nvidia-cudnn-cu12
 pip install nvidia-cuda-cupti-cu12
 ```
 
-After installing cuDNN and cuPTI, the user may install `cuquantum-python-jax` using `pip` using either:
-
-```
-pip install cuquantum-python-jax
-```
-
-or
+After installing cuDNN and cuPTI, the user may install `cuquantum-python-jax` with `pip` using either:
 
 ```
-pip install cuquantum-python-cu12[jax]
+pip install --no-build-isolation cuquantum-python-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 
+in which case the CUDA version will be detected, or one of
 
 ```
-pip install jax[cuda12]  # install cuPTI and cuDNN together with CUDA-enabled JAX
+pip install --no-build-isolation cuquantum-python-cu12[jax]
+pip install --no-build-isolation cuquantum-python-cu13[jax]
 ```
 
-and either 
-
-```
-pip install cuquantum-python-jax
-```
-
-or
-
-```
-pip install cuquantum-python-cu12[jax]
-```
+where the CUDA version is explicitly specified on cuquantum-python.
 
-Warning: if the user has an installation of CUDA outside of `pip`, this may create conflicts and undefined behavior.
+Note:
+1. If cuDNN and cuPTI are installed with `pip`, the user does not need to specify library folders in `LD_LIBRARY_PATH`.
+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.
 
 #### Installing from source
 
@@ -90,8 +75,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~=26.1.0 for CUDA 12 or cuquantum-python-cu13~=26.1.0 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.4/cuquantum/densitymat/jax/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES
+#
+# 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 import (
+    ElementaryOperator,
+    MatrixOperator,
+    OperatorTerm,
+    Operator
+)