cuquantum-python-jax-cu12 0.0.7__tar.gz → 0.0.8__tar.gz

{cuquantum_python_jax_cu12-0.0.7/cuquantum_python_jax_cu12.egg-info → cuquantum_python_jax_cu12-0.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuquantum-python-jax-cu12
-Version: 0.0.7
+Version: 0.0.8
 Summary: NVIDIA cuQuantum Python JAX
 Author-email: NVIDIA Corporation <cuquantum-python@nvidia.com>
 License-Expression: BSD-3-Clause
@@ -21,7 +21,7 @@ Requires-Python: >=3.11.0
 Description-Content-Type: text/markdown
 License-File: LICENSE
 License-File: NV.LICENSE
-Requires-Dist: cuquantum-python-cu12~=26.9.0
+Requires-Dist: cuquantum-python-cu12~=26.11.0
 Requires-Dist: jax[cuda12-local]>=0.8
 Dynamic: license-file
 
@@ -106,7 +106,7 @@ Runtime dependencies of the cuQuantum Python JAX package include:
 * Python 3.11+
 * An NVIDIA GPU with compute capability 7.5+
 * CUDA Toolkit 12.x or 13.x
-* cuquantum-python-cu12~=26.9.0 for CUDA 12 or cuquantum-python-cu13~=26.9.0 for CUDA 13
+* cuquantum-python-cu12~=26.11.0 for CUDA 12 or cuquantum-python-cu13~=26.11.0 for CUDA 13
 * jax[cuda12-local]>=0.8 for CUDA 12 or jax[cuda13-local]>=0.8 for CUDA 13
 
 ## Developer Notes

{cuquantum_python_jax_cu12-0.0.7 → cuquantum_python_jax_cu12-0.0.8}/README.md

@@ -79,7 +79,7 @@ Runtime dependencies of the cuQuantum Python JAX package include:
 * Python 3.11+
 * An NVIDIA GPU with compute capability 7.5+
 * CUDA Toolkit 12.x or 13.x
-* cuquantum-python-cu12~=26.9.0 for CUDA 12 or cuquantum-python-cu13~=26.9.0 for CUDA 13
+* cuquantum-python-cu12~=26.11.0 for CUDA 12 or cuquantum-python-cu13~=26.11.0 for CUDA 13
 * jax[cuda12-local]>=0.8 for CUDA 12 or jax[cuda13-local]>=0.8 for CUDA 13
 
 ## Developer Notes

{cuquantum_python_jax_cu12-0.0.7 → cuquantum_python_jax_cu12-0.0.8}/cuquantum/densitymat/jax/cppsrc/cudensitymat.h

@@ -64,7 +64,7 @@
 
 #define CUDENSITYMAT_MAJOR 0 //!< cuDensityMat major version.
 #define CUDENSITYMAT_MINOR 7 //!< cuDensityMat minor version.
-#define CUDENSITYMAT_PATCH 0 //!< cuDensityMat patch version.
+#define CUDENSITYMAT_PATCH 1 //!< cuDensityMat patch version.
 #define CUDENSITYMAT_VERSION (CUDENSITYMAT_MAJOR * 10000 + CUDENSITYMAT_MINOR * 100 + CUDENSITYMAT_PATCH)
 
 
@@ -3774,11 +3774,7 @@ cudensitymatStatus_t cudensitymatEigenDecompositionConfigure(
  * (e) the representative `state` is a single-site (1-mode) MPS (an MPS
  * with at least two sites is required by the DMRG engine);
  * (f) the state data type is not `CUDA_R_64F` or `CUDA_C_64F`, or
- * `computeType` is not `CUDENSITYMAT_COMPUTE_64F`; or
- * (g) shift-invert DMRG is configured with `NUM_SITES == 2` for an MPS
- * containing exactly two tensors. This last restriction is a two-site
- * ProjectionMPS dependency limitation: `NUM_SITES == 1` remains supported
- * for a two-tensor MPS.
+ * `computeType` is not `CUDENSITYMAT_COMPUTE_64F`.
  */
 cudensitymatStatus_t cudensitymatEigenDecompositionPrepare(
                     const cudensitymatHandle_t handle,
@@ -3928,6 +3924,12 @@ cudensitymatStatus_t cudensitymatDestroyWorkspace(
 /**
  * \brief Queries the required workspace buffer size.
  * 
+ * \note Two-site MPS operations can execute without caller-provided host
+ * scratch even when a nonzero host workspace size is reported. If host
+ * scratch is provided, its size must be at least the reported value.
+ * Providing the reported host workspace is recommended for performance.
+ * Device scratch requirements remain mandatory.
+ *
  * \param[in] handle Library handle.
  * \param[in] workspaceDescr Workspace descriptor.
  * \param[in] memSpace Memory space.

{cuquantum_python_jax_cu12-0.0.7 → cuquantum_python_jax_cu12-0.0.8}/cuquantum/densitymat/jax/cppsrc/cudensitymat_jax.cpp

@@ -56,7 +56,7 @@ xla::ffi::Error OperatorActionImpl(cudaStream_t stream,
             FFI_CUDA_ERROR_CHECK(cudaStreamSynchronize(stream));
         }
 
-        // Attach storage to input state.
+        // Process buffers for all inputs, including input state, elementary/matrix operators and coefficients.
         std::vector<void*> stateInComponentBufs;
         std::vector<size_t> stateInComponentSizes;
         for (int i = 0; i < inBufs.size(); ++i) {
@@ -103,6 +103,7 @@ xla::ffi::Error OperatorActionImpl(cudaStream_t stream,
             FFI_CUDA_ERROR_CHECK(cudaStreamSynchronize(stream));
         }
 
+        // Attach storage to input state.
         FFI_CUDM_ERROR_CHECK(cudensitymatStateAttachComponentStorage(handle,
                                                                      stateIn,
                                                                      numStateComponents,
@@ -147,9 +148,6 @@ xla::ffi::Error OperatorActionImpl(cudaStream_t stream,
             stateOutComponentSizes.data())
         );
 
-
-        // Execute operator action.
-
         // Attach batched coefficients if needed.
         if (operatorTermBatchedCoeffsPtrs.size() > 0 || operatorProductBatchedCoeffsPtrs.size() > 0) {
             FFI_CUDM_ERROR_CHECK(cudensitymatAttachBatchedCoefficients(
@@ -166,8 +164,7 @@ xla::ffi::Error OperatorActionImpl(cudaStream_t stream,
         // Initialize output state to zero.
         FFI_CUDM_ERROR_CHECK(cudensitymatStateInitializeZero(handle, stateOut, stream));
 
-        FFI_CUDA_ERROR_CHECK(cudaStreamSynchronize(stream));
-
+        // Execute operator action.
         FFI_CUDM_ERROR_CHECK(cudensitymatOperatorComputeAction(handle,
                                                                superoperator,
                                                                0.0, // time
@@ -193,8 +190,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(
     OperatorActionImpl,
     xla::ffi::Ffi::Bind()
         .Ctx<xla::ffi::PlatformStream<cudaStream_t>>()
-        .RemainingArgs() // all input buffers
-        .RemainingRets() // all output buffers (workspace first, then state out components)
+        .RemainingArgs() // all input buffers (state_in components, elementary/matrix operators, coefficients)
+        .RemainingRets() // all output buffers (workspace first, then state_out components)
         .Attr<xla::ffi::Span<const int64_t>>("other_in_types")
         .Attr<xla::ffi::Span<const int64_t>>("other_in_ptrs")
         .Attr<int64_t>("batch_size")
@@ -316,7 +313,6 @@ xla::ffi::Error OperatorActionBackwardDiffImpl(cudaStream_t stream,
                                                                      stateOutAdjComponentSizes.data()));
 
         // Attach storage to output state.
-        // outBufs layout: [workspace, state in adj components..., grad bufs...]
         std::vector<void*> stateInAdjComponentBufs;
         std::vector<size_t> stateInAdjComponentSizes;
 
@@ -353,10 +349,18 @@ xla::ffi::Error OperatorActionBackwardDiffImpl(cudaStream_t stream,
                                                             workspacePtrAligned,
                                                             workspaceSizeAligned));
 
-        // Initialize output state to zero.
+        // Initialize state_in_adj to zero.
         FFI_CUDM_ERROR_CHECK(cudensitymatStateInitializeZero(handle, stateInAdj, stream));
 
-        // Execute operator action backward differentiation.
+        // Zero out the persistent gradient buffers.
+        for (int i = 0; i < otherOutPtrs.size(); ++i) {
+            xla::ffi::Result<xla::ffi::AnyBuffer> resBuf =
+                outBufs.get<xla::ffi::AnyBuffer>(i + numStateComponents + 1).value();
+            FFI_CUDA_ERROR_CHECK(cudaMemsetAsync(reinterpret_cast<void*>(otherOutPtrs[i]),
+                                                 0,
+                                                 resBuf->size_bytes(),
+                                                 stream));
+        }
 
         // Attach batched coefficients if needed.
         if (operatorTermBatchedCoeffsPtrs.size() > 0 || operatorProductBatchedCoeffsPtrs.size() > 0) {
@@ -371,6 +375,7 @@ xla::ffi::Error OperatorActionBackwardDiffImpl(cudaStream_t stream,
                 operatorProductBatchedCoeffs.data()));
         }
 
+        // Execute operator action backward differentiation.
         FFI_CUDM_ERROR_CHECK(cudensitymatOperatorComputeActionBackwardDiff(handle,
                                                                            superoperator,
                                                                            0.0, // time
@@ -384,13 +389,7 @@ xla::ffi::Error OperatorActionBackwardDiffImpl(cudaStream_t stream,
                                                                            workspaceDesc,
                                                                            stream));
 
-        // The Python gradient callbacks (for both scalar and tensor operators) dispatch
-        // their += accumulation operations on CuPy's default CUDA stream, which may differ
-        // from the XLA FFI stream used here. Synchronize the default stream before copying
-        // so that all callback writes are guaranteed to be visible to the memcpy below.
-        FFI_CUDA_ERROR_CHECK(cudaStreamSynchronize(cudaStreamDefault));
-
-        // outBufs layout (after workspace and numStateComponents state adj bufs): [grad bufs].
+        // Copy gradient buffers from CuPy arrays to XLA buffers.
         for (int i = 0; i < otherOutPtrs.size(); ++i) {
             xla::ffi::Result<xla::ffi::AnyBuffer> resBuf =
                 outBufs.get<xla::ffi::AnyBuffer>(i + numStateComponents + 1).value();
@@ -418,8 +417,9 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(
     OperatorActionBackwardDiffImpl,
     xla::ffi::Ffi::Bind()
         .Ctx<xla::ffi::PlatformStream<cudaStream_t>>()
-        .RemainingArgs() // all input buffers
-        .RemainingRets() // all output buffers (workspace first, then state in adj components, then grad bufs)
+        .RemainingArgs() // all input buffers (state_in components, state_out_adj components,
+                         // elementary/matrix operators, coefficients)
+        .RemainingRets() // all output buffers (workspace first, then state_in_adj components, then grad bufs)
         .Attr<xla::ffi::Span<const int64_t>>("other_in_types")
         .Attr<xla::ffi::Span<const int64_t>>("other_in_ptrs")
         .Attr<xla::ffi::Span<const int64_t>>("other_out_ptrs")

{cuquantum_python_jax_cu12-0.0.7 → cuquantum_python_jax_cu12-0.0.8}/cuquantum/densitymat/jax/pysrc/context.py

@@ -300,46 +300,40 @@ class StateContext:
         )
         self.logger.debug(f"Created output state at {hex(self._state_out)}")
 
-        # The state adjoints are to be set in create_adjoint_buffers when backward rule is called.
+        # The state adjoints are created by the first backward rule, in create_adjoint_buffers.
         self._state_in_adj = None
         self._state_out_adj = None
 
     def create_adjoint_buffers(self):
         """
-        Create adjoint buffers for the input and output states.
+        Create the input and output state adjoints once per config; later backward passes reuse them.
 
-        Frees any previously allocated adjoint handles before allocating new
-        ones so that repeated backward passes do not leak C-side GPU memory.
+        Like the input and output handles, the adjoint handles are config-only (buffers are attached
+        on every call), and traced backward primitives embed their addresses as static parameters.
+        They must therefore live as long as this context: destroying them here would leave earlier
+        traced or compiled backward actions pointing at freed handles. free() releases them.
         """
-        if self._state_in_adj is not None:
-            cudm.destroy_state(self._state_in_adj)
-            self.logger.debug(f"Destroyed stale input state adjoint at {hex(self._state_in_adj)}")
-            self._state_in_adj = None
-
-        if self._state_out_adj is not None:
-            cudm.destroy_state(self._state_out_adj)
-            self.logger.debug(f"Destroyed stale output state adjoint at {hex(self._state_out_adj)}")
-            self._state_out_adj = None
-
-        self._state_in_adj = cudm.create_state(
-            CudensitymatContext._handle,
-            self.state_purity,
-            len(self._space_mode_extents),
-            self._space_mode_extents,
-            self.batch_size,
-            self._data_type
-        )
-        self.logger.debug(f"Created input state adjoint at {hex(self._state_in_adj)}")
-
-        self._state_out_adj = cudm.create_state(
-            CudensitymatContext._handle,
-            self.state_purity,
-            len(self._space_mode_extents),
-            self._space_mode_extents,
-            self.batch_size,
-            self._data_type
-        )
-        self.logger.debug(f"Created output state adjoint at {hex(self._state_out_adj)}")
+        if self._state_in_adj is None:
+            self._state_in_adj = cudm.create_state(
+                CudensitymatContext._handle,
+                self.state_purity,
+                len(self._space_mode_extents),
+                self._space_mode_extents,
+                self.batch_size,
+                self._data_type
+            )
+            self.logger.debug(f"Created input state adjoint at {hex(self._state_in_adj)}")
+
+        if self._state_out_adj is None:
+            self._state_out_adj = cudm.create_state(
+                CudensitymatContext._handle,
+                self.state_purity,
+                len(self._space_mode_extents),
+                self._space_mode_extents,
+                self.batch_size,
+                self._data_type
+            )
+            self.logger.debug(f"Created output state adjoint at {hex(self._state_out_adj)}")
 
     def free(self):
         """

{cuquantum_python_jax_cu12-0.0.7 → cuquantum_python_jax_cu12-0.0.8/cuquantum_python_jax_cu12.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuquantum-python-jax-cu12
-Version: 0.0.7
+Version: 0.0.8
 Summary: NVIDIA cuQuantum Python JAX
 Author-email: NVIDIA Corporation <cuquantum-python@nvidia.com>
 License-Expression: BSD-3-Clause
@@ -21,7 +21,7 @@ Requires-Python: >=3.11.0
 Description-Content-Type: text/markdown
 License-File: LICENSE
 License-File: NV.LICENSE
-Requires-Dist: cuquantum-python-cu12~=26.9.0
+Requires-Dist: cuquantum-python-cu12~=26.11.0
 Requires-Dist: jax[cuda12-local]>=0.8
 Dynamic: license-file
 
@@ -106,7 +106,7 @@ Runtime dependencies of the cuQuantum Python JAX package include:
 * Python 3.11+
 * An NVIDIA GPU with compute capability 7.5+
 * CUDA Toolkit 12.x or 13.x
-* cuquantum-python-cu12~=26.9.0 for CUDA 12 or cuquantum-python-cu13~=26.9.0 for CUDA 13
+* cuquantum-python-cu12~=26.11.0 for CUDA 12 or cuquantum-python-cu13~=26.11.0 for CUDA 13
 * jax[cuda12-local]>=0.8 for CUDA 12 or jax[cuda13-local]>=0.8 for CUDA 13
 
 ## Developer Notes

cuquantum_python_jax_cu12-0.0.8/cuquantum_python_jax_cu12.egg-info/requires.txt

@@ -0,0 +1,2 @@
+cuquantum-python-cu12~=26.11.0
+jax[cuda12-local]>=0.8

{cuquantum_python_jax_cu12-0.0.7 → cuquantum_python_jax_cu12-0.0.8}/pyproject.toml

@@ -17,7 +17,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "cuquantum-python-jax-cu12"
-version = "0.0.7"
+version = "0.0.8"
 description = "NVIDIA cuQuantum Python JAX"
 readme = "README.md"
 license = "BSD-3-Clause"
@@ -41,7 +41,7 @@ classifiers = [
     "Environment :: GPU :: NVIDIA CUDA :: 12",
 ]
 dependencies = [
-    "cuquantum-python-cu12~=26.9.0",
+    "cuquantum-python-cu12~=26.11.0",
     "jax[cuda12-local]>=0.8",
 ]
 

{cuquantum_python_jax_cu12-0.0.7 → cuquantum_python_jax_cu12-0.0.8}/pyproject.toml.template

@@ -17,7 +17,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "cuquantum-python-jax-cu@CUDA_MAJOR_VER@"
-version = "0.0.7"
+version = "0.0.8"
 description = "NVIDIA cuQuantum Python JAX"
 readme = "README.md"
 license = "BSD-3-Clause"
@@ -41,7 +41,7 @@ classifiers = [
     "@CUDA_CLASSIFIER@",
 ]
 dependencies = [
-    "cuquantum-python-cu@CUDA_MAJOR_VER@~=26.9.0",
+    "cuquantum-python-cu@CUDA_MAJOR_VER@~=26.11.0",
     "jax[cuda@CUDA_MAJOR_VER@-local]@JAX_VERSION_SPEC@",
 ]
 

cuquantum_python_jax_cu12-0.0.7/cuquantum_python_jax_cu12.egg-info/requires.txt

@@ -1,2 +0,0 @@
-cuquantum-python-cu12~=26.9.0
-jax[cuda12-local]>=0.8