Make qchem jax compatible

[austingmhuang]

Context:
Autograd deprecation means we are moving to using JAX for auto-differentiation. Since the requires_grad keyword is not supported by JAX arrays, we need a different solution.

Description of the Change:
We keep backwards compatibility to pnp by checking what interface the user is using. We check this by checking the interface of the tensor. If they are using autograd, then we stick to the old workflow and check requires_grad using getattr().

If the user inputs a jax array for any of [coordinates, coefficients, alpha], we assume that the user wants to use JAX and define all undefined coeffs/alphas using jax arrays. This means that if a user decides to mix pnp with jax, we don't hard cast the rest into either since we can't make a decision; therefore it'll result in a warning about mixing these two.

WHEN USING JAX:
If users wish to differentiate any one of these parameters they should mark the parameter they want differentiable using the JAX UI, e.g. jax.grad(..., argnums = <indice(s) of differentiable parameter(s)>)(*args). In our case, due to technical limitations, *args must always be exactly[coordinates, coefficients, alpha]. No other order is allowed and you cannot omit any of them. Unfortunately, this also includes when you are NOT using jax.grad or a jax function, like when you define diff_hamiltonian(...)(*args), the args here (if using JAX) also need to be exactly [coordinates, coefficients, alpha]. When you do decide to differentiate, doing jax.grad(..., argnums=1)(coordinates, coefficients, alpha) would mean you want coefficients to be differentiable. Note this is a departure from the UI of qml.grad, where you could do qml.grad(..., argnum=0)(coefficients) instead.

Additional notes:
UI for qml.grad and all the other stuff is unchanged for autograd and pnp users. However, if you are using JAX and trying to use the args keyword in molecular_hamiltonian and related hamiltonians, you will need to define all of [coordinates, coefficients, alpha] as well since it goes downstream to diff_hamiltonian(...)(*args).

Benefits:
Now JAX compatible.

Possible Drawbacks:
More changes may be needed to JIT, may have performance issues. Different UI for qml.grad and jax.grad. Different expectations for args keyword for jax arrays and pnp arrays.

Related GitHub Issues:
[sc-69776] [sc-69778]

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md with:

• A one-to-two sentence description of the change. You may include a small working example for new features.
• A link back to this PR.
• Your name (or GitHub username) in the contributors section.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.64%. Comparing base (c5fd5bc) to head (705395a).
Report is 2 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #6096      +/-   ##
==========================================
- Coverage   99.64%   99.64%   -0.01%     
==========================================
  Files         469      468       -1     
  Lines       44331    44065     -266     
==========================================
- Hits        44173    43907     -266     
  Misses        158      158              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[soranjh]

Thanks @austingmhuang for fixing this. Left some comments, mostly minor, but my main point is about duplicating many tests for jax, which seems rather unnecessary.

[soranjh]

Thanks @austingmhuang, please update the changelog before merging.