catalyst
catalyst copied to clipboard
Add operations for setting state and basis state and integrate this with `skip_initial_state_prep`
Context: It has been determined that qml.StatePrep takes too long to compile. To facilitate this, it was determined that the optimization of setting the initial state to the operand of qml.StatePrep could help. This is already implemented in lightning.qubit and lightning.kokkos when using them without @qjit.
Description of the Change: This change adds two new JAXPR and MLIR operations. One for qml.StatePrep and another for qml.BasisState. This change includes:
- Changes in the runtime to add support for this new operations, including changes to the QuantumDevice interface
- Changes in the TOML specification
- Changes in the frontend
- A runtime implementation of
qml.BasisStatewhich provides validation at runtime. This is needed because when using jax.numpy.arrays all values are staged.
It was necessary to add new operations instead of reusing quantum.custom op because of the signature of these operations. I.e., quantum.custom expects floating point parameters, but qml.StatePrep expects a vector of complex numbers and qml.BasisState expects an integer index.
Benefits: Possibly faster compile times (has not been verified).
Possible Drawbacks: The introduction of new operations and extension of the runtime for operations that not all devices will need can be seen as a code smell. Kokkos does not implement this since it does not implement the new device API.
Related GitHub Issues: #939
close #939
Must be merged after https://github.com/PennyLaneAI/pennylane-lightning/pull/843
TODO: update lightning git tag and version
[sc-69069]