JAX Extension APIs#

cuQuantum Python JAX is an experimental extension for cuQuantum Python, designed to provide selected functionality of the cuQuantum SDK in a JAX-compatible manner. This extension allows JAX-based frameworks to directly integrate with the cuQuantum API.

  • cuDensityMat-JAX (cuquantum.densitymat.jax) — JAX interface to the Operator Action API from cuDensityMat, for use in quantum dynamics simulation frameworks.

  • cuStabilizer-JAX (cuquantum.stabilizer.jax) — JAX interface to selected cuStabilizer routines, for use in stabilizer simulation and quantum error correction.

Release Notes#

cuQuantum Python JAX v0.0.6#

  • New features:

    • Added support for jax.jit inside jax.vmap transformations, enabling the composition jax.vmap(jax.jit(operator_action), ...).

    • Added support for VJP (vector-Jacobian product) on vmapped operators and states, enabling gradient computation through jax.vmap-wrapped operator_action() calls with respect to batched operator and state parameters.

    • Added cuquantum.stabilizer.jax.matmul_gf2_spdn(), a new API for GF(2) sparse-dense matrix multiplication C = A @ B via cuStabilizer, where A is in CSR format and B is a bit-packed dense matrix.

  • Bugs fixed:

    • Fixed an incorrect result layout in the backward pass of operator_action(), which could cause incorrect gradient values to be returned.

Compatibility notes:

  • Requires JAX version >=0.8 for both CUDA 12 and CUDA 13.

  • Python 3.14 is now supported.

cuQuantum Python JAX v0.0.5#

  • New features:

    • Gradients can now be computed with respect to any JAX-traceable parameter entering the operator or state construction without defining gradient callbacks (see example6_grad.py for an example). All callback-related arguments have been removed from the public API, and hence data buffers previously constructed dynamically via callbacks must now be pre-computed statically before being passed to operator_action(). The removed arguments are:

  • Known issues:

    • Nested jax.vmap transformations are not currently supported.

    • jax.grad transformations inside jax.vmap are not currently supported.

cuQuantum Python JAX v0.0.4#

  • New features:

    • Added support for vector-jacobian product (VJP) transformation of operator action with batched input operators and batched input states.

    • Users can now differentiate with respect to parameters implicit in the input operator or input state instead of explicitly specifying parameter gradients in gradient callbacks. See example8_gradient_attachment.py for an example.

  • Bugs fixed:

    • Fixed a bug that caused double freeing of cuDensityMat library pointers.

    • Fixed a bug that triggered an insufficient workspace runtime error when executing regular operator action after evaluating its vector-jacobian-product (VJP) function transformation.

  • Other changes:

    • Users need to pass in jax.ShapeDtypeStruct objects instead of jax.Array objects for data buffers dynamically constructed by callbacks. The affected input arguments are:

    • The diagonal offsets argument of the ElementaryOperator constructor is renamed from offsets to diag_offsets, which is used to construct multidiagonal elementary operators.

    • Modification of an Operator after it has been used in an operator action is now disabled.

    • When installing cuQuantum Python JAX, the user needs to pass the --no-build-isolation option to pip and ensure that all build dependencies are pre-installed.

cuQuantum Python JAX v0.0.3#

  • Previously, cuQuantum Python JAX set jax_enable_x64=True as a side effect on import. Now, users must set jax_enable_x64 to True before importing the cuQuantum Python JAX module.

Compatibility notes:

  • cuQuantum Python JAX now supports CUDA 13 in addition to CUDA 12.

  • cuQuantum Python JAX supports JAX version >=0.8.0 and <0.9.0 for CUDA 13

  • cuQuantum Python JAX supports JAX version >=0.5.0 and <0.7.0 for CUDA 12

cuQuantum Python JAX v0.0.2#

  • Bugs fixed:

    • Fixed an issue with packaging of the cuQuantum Python JAX extension which rendered the package uninstallable in certain situations.

cuQuantum Python JAX v0.0.1#

  • Initial release of the JAX extension exposes the Operator Action API cuquantum.densitymat.jax.operator_action() from cuDensityMat to enable integration of cuQuantum Python with JAX-based quantum dynamics simulation frameworks.

  • Known issues:

    • Multiple structurally equivalent operators, which involve callback functions for elementary/matrix operators, will result in an undefined behavior when their action on a quantum state is evaluated in consecutive calls in the same python interpreter session. As a workaround, the user may call jax.clear_caches() in between consecutive calls to operator_action() involving these operators.

    • Providing data buffers that may be overwritten by callbacks but have been initialized with the same initial values will result in an undefined behavior inside a jax.jit scope. As a workaround, the user should ensure arrays are initialized with different initial values when using them for dynamically constructed data buffers of different elementary/matrix operators.

Compatibility notes:

  • cuQuantum Python JAX supports CUDA 12.

  • cuQuantum Python JAX supports JAX version >=0.5.0 and <0.7.0.

Module APIs#