Thanks for the response @isaacdevlugt. I surely can use default.qubit but it’s performance will not scale with the number of qubits. Also, my main aim is to execute the circuit on GPUs so jax.pmap seems a nice alternative but have to look into the sharding stuff to optimize this further.
Thank you very much @isaacdevlugt.