JAX compatibility with cudaq

[ankit27kh]

JAX does not seem to be compatible with cudaq currently. It can be a great extension that will make the compute faster with functions like vmap and optimizing with optax with JIT enabled.
Are there any plans for this?

[efratshabtai]

Hi @ankit27kh, can you give us an example of what you are trying to do with JAX and CUDA-Q? What would you like to see supported in CUDA-Q?

[ankit27kh]

Hi, @efratshabtai , here's a simple example:

import cudaq
import jax
import numpy as np

@cudaq.kernel
def kernel(x: np.ndarray,params:np.ndarray):
    qubits= cudaq.qvector(x)
    h(qubits)
    rx(params[0],qubits[0])
    cry(params[1],qubits[0],qubits[1])


x=np.random.random([10,4])
params=np.random.random([2])

cost=lambda xi:cudaq.observe(kernel,cudaq.spin.z(1),xi,params,shots_count=10000).expectation()

print(cost(x[0]))                                            # this line executes 
print(jax.vmap(cost)(x))                              # this line gives an error

This is the error:

RuntimeError: error: Can not handle conversion of python type <class 'jax._src.interpreters.batching.BatchTracer'> to MLIR type.

With proper JAX support, the offending line should execute and give results for all x values at once.

We'll also need better support for different types, like multi dimensional numpy and jax.numpy arrays for the kernel and gates.

[ankit27kh]

For context, I am coming from PennyLane, where JAX is tightly integrated and helps a lot in training models quickly.