catalyst icon indicating copy to clipboard operation
catalyst copied to clipboard

Add operations for setting state and basis state and integrate this with `skip_initial_state_prep`

Open rauletorresc opened this issue 1 year ago • 9 comments

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.BasisState which 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]

rauletorresc avatar Jul 23 '24 23:07 rauletorresc