Add differentiable backend using Jax

[madagra]

Motivations

Currently, we only support PyTorch as DifferentiableBackend. However, it is useful to have a more general interface to support other backends such as Jax. The benefits of Jax over PyTorch are several:

• fully functional interface
• out-of-the-box JIT compilation of models also with complex number
• usually performance of Jax models is significantly higher than PyTorch even without compilation

Proposal

Adapt the current DifferentiableBackend interface to support Jax as a differentiable backend for the code. In essence, one should be able to write an abstract program with qadence and its execution should yield a Jax array in case that Jax differentiable backend is selected. In this way, qadence models will be compatible with any Jax-based optimization workflow and library.

To bear in mind while performing this implementation:

• the QuantumModel and generally high-level interface of the package should not change
• It is probably necessary to introduce a generic Tensor datatype that handles tensors in both Jax and PyTorch. If possible, do not proceed the Pennylane way, i.e. using the autoray package. The interface is not clean and the package is not really maintained currently.
• make sure that also the GPSR differentiation mode is taken into account. If too long, this can be handled in a separate issue.

[dominikandreasseitz]

@Roland-djee

[RolandMacDoland]

@dominikandreasseitz This can be closed right ?