cuquantum-python-jax-cu13 0.0.6__tar.gz → 0.0.7__tar.gz

{cuquantum_python_jax_cu13-0.0.6/cuquantum_python_jax_cu13.egg-info → cuquantum_python_jax_cu13-0.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuquantum-python-jax-cu13
-Version: 0.0.6
+Version: 0.0.7
 Summary: NVIDIA cuQuantum Python JAX
 Author-email: NVIDIA Corporation <cuquantum-python@nvidia.com>
 License-Expression: BSD-3-Clause
@@ -21,13 +21,13 @@ Requires-Python: >=3.11.0
 Description-Content-Type: text/markdown
 License-File: LICENSE
 License-File: NV.LICENSE
-Requires-Dist: cuquantum-python-cu13~=26.6.0
+Requires-Dist: cuquantum-python-cu13~=26.9.0
 Requires-Dist: jax[cuda13-local]>=0.8
 Dynamic: license-file
 
 # cuQuantum Python JAX
 
-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.
+cuQuantum Python JAX provides a JAX extension for cuQuantum Python. It exposes selected functionalities 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 JAX interfaces to the Operator Action API from the cuDensityMat library and the GF(2) sparse-dense matrix multiplication from the cuStabilizer library.
 
 ## Documentation
 
@@ -39,15 +39,17 @@ Please visit the [NVIDIA cuQuantum Python documentation](https://docs.nvidia.com
 
 The build-time dependencies of the cuQuantum Python JAX package include:
 
+* Python 3.11+
 * CUDA Toolkit 12.x or 13.x
-* jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8,<0.9 for CUDA 13
+* cmake>=3.22
+* jax[cuda12-local]>=0.8 for CUDA 12 or jax[cuda13-local]>=0.8 for CUDA 13
 * pybind11
 * wheel
 * setuptools>=77.0.3
 
 Note:
-- cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13.
-- cuQuantum Python JAX wheels are CUDA-versioned: `cuquantum-python-jax-cu12` for CUDA 12 and `cuquantum-python-jax-cu13` for CUDA 13.
+- cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13 and the distributions are CUDA-versioned: `cuquantum-python-jax-cu12` for CUDA 12 and `cuquantum-python-jax-cu13` for CUDA 13.
+- **Build isolation note:** When installing with build isolation (`pip`'s default mode), the JAX version used at build time is determined by the build dependency specification in `pyproject.toml` and may differ from the JAX version pre-installed in the environment. Mismatched build-time and runtime JAX versions can lead to errors when importing cuQuantum's JAX extension. We recommend passing `--no-build-isolation` to `pip install` and pre-installing all build dependencies manually.
 
 #### Installation using `jax[cudaXX-local]`
 
@@ -101,10 +103,11 @@ The CUDA version is detected automatically from `$CUDA_PATH` and the wheel will
 
 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.3.0 for CUDA 12 or cuquantum-python-cu13~=26.3.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
+* cuquantum-python-cu12~=26.9.0 for CUDA 12 or cuquantum-python-cu13~=26.9.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_cu13-0.0.6 → cuquantum_python_jax_cu13-0.0.7}/README.md

@@ -1,6 +1,6 @@
 # cuQuantum Python JAX
 
-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.
+cuQuantum Python JAX provides a JAX extension for cuQuantum Python. It exposes selected functionalities 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 JAX interfaces to the Operator Action API from the cuDensityMat library and the GF(2) sparse-dense matrix multiplication from the cuStabilizer library.
 
 ## Documentation
 
@@ -12,15 +12,17 @@ Please visit the [NVIDIA cuQuantum Python documentation](https://docs.nvidia.com
 
 The build-time dependencies of the cuQuantum Python JAX package include:
 
+* Python 3.11+
 * CUDA Toolkit 12.x or 13.x
-* jax[cuda12-local]>=0.5,<0.7 for CUDA 12 or jax[cuda13-local]>=0.8,<0.9 for CUDA 13
+* cmake>=3.22
+* jax[cuda12-local]>=0.8 for CUDA 12 or jax[cuda13-local]>=0.8 for CUDA 13
 * pybind11
 * wheel
 * setuptools>=77.0.3
 
 Note:
-- cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13.
-- cuQuantum Python JAX wheels are CUDA-versioned: `cuquantum-python-jax-cu12` for CUDA 12 and `cuquantum-python-jax-cu13` for CUDA 13.
+- cuQuantum Python JAX is only supported with CUDA 12 and CUDA 13 and the distributions are CUDA-versioned: `cuquantum-python-jax-cu12` for CUDA 12 and `cuquantum-python-jax-cu13` for CUDA 13.
+- **Build isolation note:** When installing with build isolation (`pip`'s default mode), the JAX version used at build time is determined by the build dependency specification in `pyproject.toml` and may differ from the JAX version pre-installed in the environment. Mismatched build-time and runtime JAX versions can lead to errors when importing cuQuantum's JAX extension. We recommend passing `--no-build-isolation` to `pip install` and pre-installing all build dependencies manually.
 
 #### Installation using `jax[cudaXX-local]`
 
@@ -74,10 +76,11 @@ The CUDA version is detected automatically from `$CUDA_PATH` and the wheel will
 
 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.3.0 for CUDA 12 or cuquantum-python-cu13~=26.3.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
+* cuquantum-python-cu12~=26.9.0 for CUDA 12 or cuquantum-python-cu13~=26.9.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_cu13-0.0.6 → cuquantum_python_jax_cu13-0.0.7}/cuquantum/densitymat/jax/__init__.py

@@ -2,11 +2,6 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-from cuquantum.bindings._internal import cudensitymat as _cudm
-_cudm._inspect_function_pointers()  # for loading libcudensitymat.so
-
-import jax
-
 try:
     from ._build_info import check_jax_abi as _check_jax_abi
 except ImportError:
@@ -15,13 +10,11 @@ else:
     _check_jax_abi()
     del _check_jax_abi
 
-if not jax.config.jax_enable_x64:
-    raise RuntimeError("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
+    Operator,
+    SimplifierConfig
 )

{cuquantum_python_jax_cu13-0.0.6 → cuquantum_python_jax_cu13-0.0.7}/cuquantum/densitymat/jax/cppsrc/cudensitymat.h

@@ -63,7 +63,7 @@
 // LIBRARY VERSION
 
 #define CUDENSITYMAT_MAJOR 0 //!< cuDensityMat major version.
-#define CUDENSITYMAT_MINOR 6 //!< cuDensityMat minor version.
+#define CUDENSITYMAT_MINOR 7 //!< cuDensityMat minor version.
 #define CUDENSITYMAT_PATCH 0 //!< cuDensityMat patch version.
 #define CUDENSITYMAT_VERSION (CUDENSITYMAT_MAJOR * 10000 + CUDENSITYMAT_MINOR * 100 + CUDENSITYMAT_PATCH)
 
@@ -282,7 +282,8 @@ typedef enum
  */
 typedef enum
 {
-  CUDENSITYMAT_EIGEN_SCOPE_SPLIT_DMRG = 0, ///< DMRG-based split decomposition (default)
+  CUDENSITYMAT_EIGEN_SCOPE_SPLIT_DMRG = 0,              ///< DMRG-based split decomposition (default)
+  CUDENSITYMAT_EIGEN_SCOPE_SPLIT_SHIFT_INVERT_DMRG = 1, ///< Shift-invert DMRG split decomposition (eigenpair nearest a target energy)
 } cudensitymatEigenDecompositionScopeSplitKind_t;
 
 /**
@@ -298,8 +299,9 @@ typedef enum
  */
 typedef enum
 {
-  CUDENSITYMAT_EIGEN_APPROACH_KRYLOV = 0, ///< Krylov subspace method
-//CUDENSITYMAT_EIGEN_APPROACH_DAVIDSON = 1, ///< Davidson method
+  CUDENSITYMAT_EIGEN_APPROACH_KRYLOV = 0, ///< Krylov subspace (block-Lanczos) eigensolver
+  CUDENSITYMAT_EIGEN_APPROACH_LINEAR = 1, ///< Iterative linear solver (currently MINRES)
+//CUDENSITYMAT_EIGEN_APPROACH_DAVIDSON = 2, ///< Davidson method
 } cudensitymatEigenDecompositionApproachKind_t;
 
 /**
@@ -322,8 +324,9 @@ typedef enum
   CUDENSITYMAT_EIGEN_SPLIT_SCOPE_KIND = 0,           ///< int32_t (cudensitymatEigenDecompositionScopeSplitKind_t): Split kind
 //CUDENSITYMAT_EIGEN_FULL_SCOPE_KIND = 1,            ///< int32_t (cudensitymatEigenDecompositionScopeFullKind_t): Full-scope kind
 //CUDENSITYMAT_EIGEN_FULL_SCOPE_EXACT_CONFIG = 2,    ///< cudensitymatEigenDecompositionScopeFullExactConfig_t: Full exact-scope configuration
-  CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_CONFIG = 3,    ///< cudensitymatEigenDecompositionScopeSplitDMRGConfig_t: DMRG split configuration
+  CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_CONFIG = 3,    ///< cudensitymatEigenDecompositionScopeSplitDMRGConfig_t: DMRG split configuration (shared by the DMRG and shift-invert DMRG split kinds)
   CUDENSITYMAT_EIGEN_APPROACH_KRYLOV_CONFIG = 10,    ///< cudensitymatEigenDecompositionApproachKrylovConfig_t: Krylov approach configuration
+  CUDENSITYMAT_EIGEN_APPROACH_LINEAR_CONFIG = 11,    ///< cudensitymatEigenDecompositionApproachLinearConfig_t: Linear-solver approach configuration
 } cudensitymatEigenDecompositionAttribute_t;
 
 // /**
@@ -378,12 +381,21 @@ typedef enum
  */
 typedef enum
 {
-  CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_NUM_SITES = 0,                       ///< int32_t: Number of sites in the DMRG chain (default: 1)
-  CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_SVD_CONFIG = 1,                      ///< cudensitymatSVDConfig_t: SVD truncation configuration
-  CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_MAX_SWEEPS = 2,                      ///< int32_t: Maximum number of full L-R-L sweeps (default: 20)
-  CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_ENERGY_TOLERANCE = 3,                ///< double: Convergence threshold on the change in variational energy between consecutive sweeps; the sweep loop terminates when |E_k - E_{k-1}| falls below this value (default: 1e-10)
+  CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_NUM_SITES = 0,                       ///< int32_t: Number of neighboring MPS sites updated in each local solve (1 or 2; default: 1). SetAttribute returns CUDENSITYMAT_STATUS_INVALID_VALUE for values <= 0 and CUDENSITYMAT_STATUS_NOT_SUPPORTED for values > 2.
+  CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_SVD_CONFIG = 1,                      ///< cudensitymatSVDConfig_t: SVD truncation configuration for 2-site updates
+  CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_MAX_SWEEPS = 2,                      ///< int32_t: Maximum number of full L-R-L sweeps (default: 20); for shift-invert, the fitting-sweep cap for each inverse application
+  CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_ENERGY_TOLERANCE = 3,                ///< double: Convergence threshold on the change in energy between consecutive sweeps for ground-state DMRG or outer power iterations for shift-invert DMRG (default: 1e-10)
+  CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_MAX_POWER_ITERATIONS = 4,            ///< int32_t: Maximum number of outer power iterations on (H - sigma)^{-1} for the shift-invert DMRG split kind; ignored by the ground-state DMRG split kind (default: 10)
 } cudensitymatEigenDecompositionScopeSplitDMRGConfigAttribute_t;
 
+/**
+ * \brief Configuration attributes for the iterative linear solver (currently MINRES).
+ */
+typedef enum
+{
+  CUDENSITYMAT_EIGEN_APPROACH_LINEAR_MAX_ITERATIONS = 0, ///< int32_t: Maximum number of iterations for the local linear solve (default: 200). The per-site solve convergence tolerance is supplied at Compute time via the in/out tolerances[] argument, not as a config attribute.
+} cudensitymatEigenDecompositionApproachLinearConfigAttribute_t;
+
 /**
  * \brief State-fitting scope (full vs split fitting).
  */
@@ -500,7 +512,7 @@ typedef enum
 typedef enum
 {
   CUDENSITYMAT_EIGEN_SPECTRUM_LARGEST = 0,       ///< Compute the largest by magnitude eigen-values of the operator
-  CUDENSITYMAT_EIGEN_SPECTRUM_SMALLEST = 1,      ///< Compute the smallest by magnitude eigen-values of the operator
+  CUDENSITYMAT_EIGEN_SPECTRUM_SMALLEST = 1,      ///< Compute the smallest by magnitude eigenvalues; for shift-invert DMRG, select the smallest |E - sigma| using the Compute target
   CUDENSITYMAT_EIGEN_SPECTRUM_LARGEST_REAL = 2,  ///< Compute the largest by the real part eigen-values of the operator
   CUDENSITYMAT_EIGEN_SPECTRUM_SMALLEST_REAL = 3, ///< Compute the smallest by the real part eigen-values of the operator
 } cudensitymatEigenDecompositionSpectrumKind_t;
@@ -685,6 +697,14 @@ typedef void * cudensitymatTimePropagationScopeSplitTDVPConfig_t;
  */
 typedef void * cudensitymatEigenDecompositionScopeSplitDMRGConfig_t;
 
+/**
+ * \brief Opaque data structure holding the iterative linear-solver configuration.
+ *
+ * \details This configuration object stores the iteration cap for the iterative
+ * linear solver; the solve tolerance is supplied to Compute.
+ */
+typedef void * cudensitymatEigenDecompositionApproachLinearConfig_t;
+
 /**
  * \brief Opaque data structure holding variational ALS configuration
  * for split-scope state fitting on a factorized (e.g., MPS) quantum state.
@@ -1631,7 +1651,8 @@ cudensitymatStatus_t cudensitymatStateComputeInnerProduct(
  *     The full non-zero diagonals are stored in a concatenated form,
  *     following the order how they appear in the `diagonalOffsets` argument.
  *     The length of each stored diagonal is equal to the full matrix dimension,
- *     padded with trailing zeros for non-main diagonals.
+ *     padded with trailing zeros for non-main diagonals. Each non-zero diagonal
+ *     is stored exactly once, hence all diagonal offsets must be unique.
  *
  * \note Currently the multi-diagonal storage format is only supported
  * by 1-body elementary tensor operators (restriction subject to lifting in future).
@@ -1648,7 +1669,9 @@ cudensitymatStatus_t cudensitymatStateComputeInnerProduct(
  * \param[in] diagonalOffsets For multi-diagonal tensor operator matrices, these are
  * the offsets of the non-zero diagonals (for example, the main diagonal has offset 0,
  * the diagonal right above the main diagonal has offset +1, the diagonal right below
- * the main diagonal has offset -1, and so on).
+ * the main diagonal has offset -1, and so on). The diagonal offsets must be unique
+ * (each non-zero diagonal is specified and stored exactly once); repeated offsets
+ * are rejected.
  * \param[in] dataType Tensor operator data type.
  * \param[in] tensorData GPU-accessible pointer to the tensor operator elements storage.
  * \param[in] tensorCallback Optional user-defined tensor callback function
@@ -1703,7 +1726,8 @@ cudensitymatStatus_t cudensitymatCreateElementaryOperator(
  *     The full non-zero diagonals are stored in a concatenated form,
  *     following the order how they appear in the `diagonalOffsets` argument.
  *     The length of each stored diagonal is equal to the full matrix dimension,
- *     padded with trailing zeros for non-main diagonals.
+ *     padded with trailing zeros for non-main diagonals. Each non-zero diagonal
+ *     is stored exactly once, hence all diagonal offsets must be unique.
  *
  * \note Currently the multi-diagonal storage format is only supported
  * by 1-body elementary tensor operators (restriction subject to lifting in future).
@@ -1724,6 +1748,8 @@ cudensitymatStatus_t cudensitymatCreateElementaryOperator(
  * \param[in] diagonalOffsets Offsets of the non-zero diagonals (for example,
  * the main diagonal has offset 0, the diagonal right above the main diagonal
  * has offset +1, the diagonal right below the main diagonal has offset -1, and so on).
+ * The diagonal offsets must be unique (each non-zero diagonal is specified and
+ * stored exactly once); repeated offsets are rejected.
  * \param[in] dataType Tensor operator data type.
  * \param[in] tensorData GPU-accessible pointer to the tensor operator elements storage,
  * where all elementary tensor operators within the batch are stored contiguously in memory.
@@ -3011,9 +3037,13 @@ cudensitymatStatus_t cudensitymatOperatorSpectrumPrepare(
 * of shape [numEigenStates, batchSize]) in GPU-accessible RAM (same data type
 * as used by the quantum state and operator).
 * \param[inout] tolerances Pointer to an F-order array of shape [numEigenStates, batchSize]
-* in CPU-accessible RAM. The initial values represent the desirable convergence tolerances
-* for all eigen-states. The returned values represent the actually achieved residual norms
-* for all eigen-states.
+* in CPU-accessible RAM. On input, the desired solver convergence tolerances for all
+* eigen-states; on output, the actually achieved solver convergence residual norms (which
+* may differ from the requested values in either direction). These characterize solver
+* convergence, not the representation error of the returned state. For the full-state solver
+* they coincide with the eigenpair residual ||H x - E x||; for split (DMRG / shift-invert)
+* solvers the returned value is the local per-site solver residual, NOT the global
+* ||H|psi> - E|psi>||. The representation error is a separate quantity (dedicated query, when available).
 * \param[in] workspace Allocated workspace descriptor.
 * \param[in] stream CUDA stream.
 * \return cudensitymatStatus_t
@@ -3572,6 +3602,71 @@ cudensitymatStatus_t cudensitymatEigenDecompositionApproachKrylovConfigGetAttrib
                     void * attributeValue,
                     size_t attributeSize);
 
+// ============================================================================
+// Linear-Solver Configuration API (eigensolver)
+// ============================================================================
+
+/**
+ * \brief Creates an eigen-decomposition iterative linear-solver approach
+ * configuration object with default settings.
+ *
+ * \param[in] handle Library handle.
+ * \param[out] config Linear-solver configuration object.
+ * \return cudensitymatStatus_t
+ */
+cudensitymatStatus_t cudensitymatCreateEigenDecompositionApproachLinearConfig(
+                    const cudensitymatHandle_t handle,
+                    cudensitymatEigenDecompositionApproachLinearConfig_t * config);
+
+/**
+ * \brief Destroys an eigen-decomposition iterative linear-solver approach
+ * configuration object.
+ *
+ * \param[in] config Linear-solver configuration object.
+ * \return cudensitymatStatus_t
+ */
+cudensitymatStatus_t cudensitymatDestroyEigenDecompositionApproachLinearConfig(cudensitymatEigenDecompositionApproachLinearConfig_t config);
+
+/**
+ * \brief Sets an attribute of the eigen-decomposition iterative linear-solver
+ * approach configuration.
+ *
+ * \param[in] handle Library handle.
+ * \param[in] config Linear-solver configuration object.
+ * \param[in] attribute Attribute to set.
+ * \param[in] attributeValue Pointer to the attribute value.
+ * \param[in] attributeSize Size of the attribute value in bytes.
+ * \return cudensitymatStatus_t
+ *
+ * \note The library captures the attribute by value at this call; the caller
+ * retains ownership of `attributeValue` and may destroy the source as soon as
+ * the call returns.
+ */
+cudensitymatStatus_t cudensitymatEigenDecompositionApproachLinearConfigSetAttribute(
+                    const cudensitymatHandle_t handle,
+                    cudensitymatEigenDecompositionApproachLinearConfig_t config,
+                    cudensitymatEigenDecompositionApproachLinearConfigAttribute_t attribute,
+                    const void * attributeValue,
+                    size_t attributeSize);
+
+/**
+ * \brief Gets an attribute of the eigen-decomposition iterative linear-solver
+ * approach configuration.
+ *
+ * \param[in] handle Library handle.
+ * \param[in] config Linear-solver configuration object.
+ * \param[in] attribute Attribute to get.
+ * \param[out] attributeValue Pointer to store the attribute value.
+ * \param[in] attributeSize Size of the buffer in bytes.
+ * \return cudensitymatStatus_t
+ */
+cudensitymatStatus_t cudensitymatEigenDecompositionApproachLinearConfigGetAttribute(
+                    const cudensitymatHandle_t handle,
+                    const cudensitymatEigenDecompositionApproachLinearConfig_t config,
+                    cudensitymatEigenDecompositionApproachLinearConfigAttribute_t attribute,
+                    void * attributeValue,
+                    size_t attributeSize);
+
 // ============================================================================
 // Eigen Decomposition API
 // ============================================================================
@@ -3584,7 +3679,8 @@ cudensitymatStatus_t cudensitymatEigenDecompositionApproachKrylovConfigGetAttrib
  * \param[in] isHermitian Specifies whether the operator is Hermitian (!=0) or not (0).
  * \param[in] spectrumKind Requested kind of the eigen-spectrum computation.
  * \param[in] scopeKind Requested decomposition scope (full vs split).
- * \param[in] approachKind Requested decomposition approach (e.g., Krylov).
+ * \param[in] approachKind Requested decomposition approach (e.g., Krylov or
+ * iterative linear solve).
  * \param[out] eigenDecomposition Eigen-decomposition computation object.
  * \return cudensitymatStatus_t
  *
@@ -3595,6 +3691,10 @@ cudensitymatStatus_t cudensitymatEigenDecompositionApproachKrylovConfigGetAttrib
  * `CUDENSITYMAT_STATUS_NOT_SUPPORTED`; only Hermitian operators are
  * supported in this release. The library trusts the caller's flag and does
  * not verify Hermiticity of the operator.
+ *
+ * \note In this release, an operator containing more than one MPO product,
+ * or an MPO product containing more than one MPO, is rejected at Create with
+ * `CUDENSITYMAT_STATUS_NOT_SUPPORTED`.
  */
 cudensitymatStatus_t cudensitymatCreateEigenDecomposition(
                     const cudensitymatHandle_t handle,
@@ -3655,17 +3755,30 @@ cudensitymatStatus_t cudensitymatEigenDecompositionConfigure(
  * \param[in] stream CUDA stream.
  * \return cudensitymatStatus_t
  *
+ * \note Prepare is transactional with respect to the cached eigen plan and
+ * workspace requirements. If Prepare fails without an intervening successful
+ * Configure, the last successfully prepared plan and workspace requirements
+ * remain valid provided the generator, state geometry, and Context distribution
+ * are unchanged. A successful Configure invalidates the prepared lifecycle, so
+ * a subsequent failed Prepare leaves the object unprepared.
+ *
  * \note In this release, Prepare returns `CUDENSITYMAT_STATUS_NOT_SUPPORTED`
  * when any of the following conditions is met:
  * (a) the `scopeKind` provided at Create is `CUDENSITYMAT_EIGEN_SCOPE_FULL`;
  * (b) the bound DMRG configuration has `CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_NUM_SITES`
  * not in `{1, 2}` (only 1-site and 2-site DMRG are supported in this release);
  * (c) `maxEigenStates != 1` (only a single eigen-pair is supported in this release);
- * (d) the `spectrumKind` provided at Create is not
- * `CUDENSITYMAT_EIGEN_SPECTRUM_SMALLEST_REAL` (only the smallest-real
- * eigenvalue is supported by the DMRG engine in this release);
+ * (d) the selected split/approach/spectrum combination is neither
+ * `SPLIT_DMRG + KRYLOV + SMALLEST_REAL` nor
+ * `SPLIT_SHIFT_INVERT_DMRG + LINEAR + SMALLEST`;
  * (e) the representative `state` is a single-site (1-mode) MPS (an MPS
- * with at least two sites is required by the DMRG engine).
+ * 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.
  */
 cudensitymatStatus_t cudensitymatEigenDecompositionPrepare(
                     const cudensitymatHandle_t handle,
@@ -3697,13 +3810,25 @@ cudensitymatStatus_t cudensitymatEigenDecompositionPrepare(
  * \param[inout] eigenstates Quantum eigenstates.
  * The initial values of the provided quantum states will be used as the
  * initial guesses for the iterative solver.
- * \param[out] eigenvalues Pointer to the eigenvalues storage (F-order array
+ * \param[inout] eigenvalues Pointer to the eigenvalues storage (F-order array
  * of shape [numEigenStates, batchSize]) in GPU-accessible RAM (same data type
- * as used by the quantum state and operator).
+ * as used by the quantum state and operator). For the shift-invert DMRG split kind,
+ * the input value is the target energy sigma (the returned eigenpair is the one whose
+ * eigenvalue is nearest sigma); it is ignored on input by other split kinds/approaches.
+ * On output, holds the computed eigenvalue(s).
  * \param[inout] tolerances Pointer to an F-order array of shape [numEigenStates, batchSize]
- * in CPU-accessible RAM. The initial values represent the desirable convergence tolerances
- * for all eigen-states. The returned values represent the actually achieved residual norms
- * for all eigen-states.
+ * in CPU-accessible RAM. On input, the desired solver convergence tolerances for all
+ * eigen-states; on output, the actually achieved solver convergence residual norms (which
+ * may differ from the requested values in either direction). These characterize solver
+ * convergence, not the representation error of the returned state. For the full-state solver
+ * they coincide with the eigenpair residual ||H x - E x||; for split (DMRG / shift-invert)
+ * solvers the returned value is the local per-site solver residual, NOT the global
+ * ||H|psi> - E|psi>||. The representation error is a separate quantity (dedicated query, when available).
+ * For shift-invert DMRG specifically, the input value is the requested per-site
+ * relative normal-equation tolerance. On output it is the maximum achieved absolute
+ * post-insertion (and, for 2-site updates, post-truncation) local normal-equation
+ * residual over every region in the final complete fitting sweep of the final outer
+ * iteration. It is not the achieved outer energy-difference metric.
  * \param[in] workspace Allocated workspace descriptor.
  * \param[in] stream CUDA stream.
  * \return cudensitymatStatus_t
@@ -3714,10 +3839,25 @@ cudensitymatStatus_t cudensitymatEigenDecompositionPrepare(
  * \note `numEigenStates != 1` returns `CUDENSITYMAT_STATUS_NOT_SUPPORTED`
  * in this release; only a single eigen-pair is supported.
  *
+ * \note `batchSize < 1` returns `CUDENSITYMAT_STATUS_INVALID_VALUE`;
+ * `batchSize > 1` returns `CUDENSITYMAT_STATUS_NOT_SUPPORTED`.
+ *
+ * \note `handle` and its distributed configuration must match the last
+ * successful Prepare. After changing either, call Prepare again before Compute.
+ *
  * \note The eigenvalue data type matches the state's data type (i.e., the
  * complex type used for the input `eigenstates`). For Hermitian operators
  * the real eigenvalue is written into the complex slot with `imag = 0`.
  *
+ * \note Shift-invert convergence requires finite real sigma, d1 > 0, and
+ * d1 < d2, where d1 and d2 are the nearest distinct-eigenvalue distances.
+ * A degenerate nearest eigenspace is allowed; exact and equidistant targets
+ * are not detected or detuned.
+ *
+ * \note Cap exhaustion returns `CUDENSITYMAT_STATUS_SUCCESS` with the final
+ * normalized approximation, Rayleigh quotient, and final-sweep local residual;
+ * these outputs do not certify outer or global convergence.
+ *
  * \note The input `eigenstates` are used as initial guesses for the iterative
  * solver. The library does not canonicalize them; the caller is responsible
  * for providing a linearly-independent (and ideally well-conditioned) set.

{cuquantum_python_jax_cu13-0.0.6 → cuquantum_python_jax_cu13-0.0.7}/cuquantum/densitymat/jax/operator_action.py

@@ -8,8 +8,6 @@ from collections.abc import Sequence
 
 import jax
 
-from cuquantum.lib.cudensitymat_jax import InputType
-
 from .pysrc.context import CudensitymatContext
 from .pysrc.operator import Operator
 from .pysrc.operator_action_prim import (
@@ -24,14 +22,33 @@ from .utils import (
     check_and_return_final_batch_size,
     check_and_return_device,
     get_original_shape,
-    get_vmap_depth,
-    is_grad_inside_vmap,
+    is_vmap_traced,
 )
 
 
 logger = logging.getLogger("cudensitymat-jax.operator_action")
 
 
+_ffi_registered = False
+_InputType = None
+
+
+def _register_ffi_targets():
+    global _ffi_registered, _InputType
+    if _ffi_registered:
+        return
+
+    from cuquantum.bindings._internal import cudensitymat as _cudm
+    _cudm._inspect_function_pointers()  # for loading libcudensitymat.so
+
+    from cuquantum.lib import cudensitymat_jax
+    _InputType = cudensitymat_jax.InputType
+
+    for _name, _value in cudensitymat_jax.registrations().items():
+        jax.ffi.register_ffi_target(_name, _value, platform="CUDA")
+    _ffi_registered = True
+
+
 def operator_action(op: Operator,
                     state_in_bufs: jax.Array | Sequence[jax.Array],
                     device: jax.Device | None = None,
@@ -49,20 +66,17 @@ def operator_action(op: Operator,
     """
     logger.info("Calling operator_action")
 
+    _register_ffi_targets()
+
+    if not jax.config.jax_enable_x64:
+        raise RuntimeError("jax_enable_x64 must be set to True to use cuQuantum Python JAX")
+
     # Process input arguments.
     if isinstance(state_in_bufs, jax.Array):
         state_in_bufs = (state_in_bufs,)
     else:
         state_in_bufs = tuple(state_in_bufs)
 
-    # Guard against nested vmap transformations, which are not supported.
-    if get_vmap_depth(state_in_bufs[0]) > 1:
-        raise NotImplementedError("operator_action does not support nested vmap transformations.")
-
-    # Guard against grad applied inside vmap, which is not supported.
-    if is_grad_inside_vmap(state_in_bufs[0]):
-        raise NotImplementedError("operator_action does not support grad transformations inside vmap.")
-
     # Check and set device from op and state.
     device = check_and_return_device(op, state_in_bufs)
 
@@ -74,12 +88,23 @@ def operator_action(op: Operator,
         device = devices[0]
         logger.info("No device specified, using the first GPU device.")
 
-    # Check state shape and maybe expand to a leading batch dimension.
-    state_batch_size, purity = get_state_batch_size_and_purity(state_in_bufs, len(op.dims))
-    batch_size = check_and_return_final_batch_size(state_in_bufs, state_batch_size, op.batch_size)
+    if op._ptr is None:
+        op._simplify()
+    op._update_metadata()  # update batch sizes
 
-    state_in_bufs, did_expand = maybe_expand_dim(state_in_bufs, len(op.dims))
-    state_shape = get_original_shape(state_in_bufs[0])
+    # Check state shape and maybe expand to a leading batch dimension.
+    state_batch_size, purity, has_explicit_batch = get_state_batch_size_and_purity(
+        state_in_bufs, op.dims)
+    batch_size = check_and_return_final_batch_size(state_batch_size, op._batch_size)
+
+    state_in_bufs, did_expand = maybe_expand_dim(state_in_bufs, has_explicit_batch)
+    # For nested vmap, get_original_shape would return a multi-dim batch prefix
+    # (outer_B, inner_B, *physical), which breaks maybe_create_state_context's
+    # single-leading-batch assumption. Fuse all batch levels into one leading dim.
+    if is_vmap_traced(state_in_bufs[0]):
+        state_shape = (batch_size, *state_in_bufs[0].shape)
+    else:
+        state_shape = tuple(state_in_bufs[0].shape)  # already has batch from maybe_expand_dim
 
     # Prepare library context for forward operator action.
     # NOTE: Assuming a single state component.
@@ -117,19 +142,19 @@ def operator_action(op: Operator,
         is_op_term_coeff_batched = op._op_term_batch_sizes[i] != 1
         if is_op_term_coeff_batched:
             dynamic_ptr = op_term_total_coeffs_ptr
-            dynamic_type = InputType.OPERATOR_TERM_BATCHED_COEFFS.value
+            dynamic_type = _InputType.OPERATOR_TERM_BATCHED_COEFFS.value
             if dynamic_ptr == 0:
                 raise RuntimeError("Missing total coefficient pointer for batched operator term coefficient.")
         else:
             dynamic_ptr = op_term_coeff_ptr
-            dynamic_type = InputType.NON_BATCHED_COEFFS.value
+            dynamic_type = _InputType.NON_BATCHED_COEFFS.value
 
         if dynamic_ptr not in op_term_coeff_metadata.ptrs:
             op_term_coeff_metadata.indices.append(i)
             op_term_coeff_metadata.types.append(dynamic_type)
             op_term_coeff_metadata.ptrs.append(dynamic_ptr)
 
-        if op_term_coeff_grad_ptr != 0:
+        if op_term_coeff_grad_ptr != 0 and op_term_coeff_grad_ptr not in op_term_coeff_grad_metadata.ptrs:
             op_term_coeff_grad_metadata.indices.append(i)
             op_term_coeff_grad_metadata.ptrs.append(op_term_coeff_grad_ptr)
             # The gradient buffer is always sized (batch_size,) regardless of whether the
@@ -155,19 +180,19 @@ def operator_action(op: Operator,
             is_op_prod_coeff_batched = op_term._op_prod_batch_sizes[j] != 1
             if is_op_prod_coeff_batched:
                 dynamic_ptr = op_prod_total_coeffs_ptr
-                dynamic_type = InputType.OPERATOR_PRODUCT_BATCHED_COEFFS.value
+                dynamic_type = _InputType.OPERATOR_PRODUCT_BATCHED_COEFFS.value
                 if dynamic_ptr == 0:
                     raise RuntimeError("Missing total coefficient pointer for batched operator product coefficient.")
             else:
                 dynamic_ptr = op_prod_coeff_ptr
-                dynamic_type = InputType.NON_BATCHED_COEFFS.value
+                dynamic_type = _InputType.NON_BATCHED_COEFFS.value
 
             if dynamic_ptr not in op_prod_coeff_metadata.ptrs:
                 op_prod_coeff_metadata.indices.append((i, j))
                 op_prod_coeff_metadata.types.append(dynamic_type)
                 op_prod_coeff_metadata.ptrs.append(dynamic_ptr)
 
-            if op_prod_coeff_grad_ptr != 0:
+            if op_prod_coeff_grad_ptr != 0 and op_prod_coeff_grad_ptr not in op_prod_coeff_grad_metadata.ptrs:
                 op_prod_coeff_grad_metadata.indices.append((i, j))
                 op_prod_coeff_grad_metadata.ptrs.append(op_prod_coeff_grad_ptr)
                 # The gradient buffer is always sized (batch_size,) regardless of whether the
@@ -182,8 +207,8 @@ def operator_action(op: Operator,
                 if base_op._ptr is not None and base_op._ptr not in base_op_metadata.ptrs:
                     base_op_metadata.indices.append((i, j, k))
                     base_op_metadata.types.append(
-                        InputType.ELEMENTARY_OPERATOR.value if base_op._is_elementary
-                        else InputType.MATRIX_OPERATOR.value
+                        _InputType.ELEMENTARY_OPERATOR.value if base_op._is_elementary
+                        else _InputType.MATRIX_OPERATOR.value
                     )
                     base_op_metadata.ptrs.append(base_op._ptr)
 
@@ -191,12 +216,18 @@ def operator_action(op: Operator,
                     base_op_grad_metadata.indices.append((i, j, k))
                     base_op_grad_metadata.ptrs.append(base_op._grad_ptr)
                     data_shape = get_original_shape(base_op.data)
-                    if base_op.batch_size == 1:
-                        # Non-batched base op: cudensitymat writes batch_size gradient tensors
+                    if is_vmap_traced(base_op.data) or base_op._batch_size > 1:
+                        # data_shape already includes the batch dim: either fused in by
+                        # get_original_shape for a vmap trace (at any extent, including a
+                        # size-1 vmap axis), or baked into the concrete array's own shape
+                        # for materialized (non-vmap) batching, where _batch_size > 1 is a
+                        # rank-verified presence signal (elementary_operator.py's
+                        # `ndim % 2 == 1` check), not an extent guess.
+                        grad_shape = data_shape
+                    else:
+                        # Genuinely unbatched: cudensitymat writes batch_size gradient tensors
                         # with batch as the last dimension, so buffer has (batch_size, *data_shape).
                         grad_shape = (batch_size, *data_shape)
-                    else:
-                        grad_shape = data_shape  # already includes batch dim
                     shape_dtype = jax.ShapeDtypeStruct(grad_shape, base_op.data.dtype)
                     base_op_grad_metadata.shape_dtypes.append(shape_dtype)
 
@@ -228,7 +259,7 @@ def operator_action(op: Operator,
     )
 
     # Undo the leading batch dim when it was added by maybe_expand_dim (single-state, non-vmap).
-    state_out_bufs = maybe_squeeze_dim(state_out_bufs, len(op.dims), did_expand)
+    state_out_bufs = maybe_squeeze_dim(state_out_bufs, did_expand)
 
     # Process output argument.
     if len(state_out_bufs) == 1: