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.70%. Comparing base (1a143d2) to head (4925338).

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #6096   +/-   ##
=======================================
  Coverage   99.70%   99.70%           
=======================================
  Files         444      444           
  Lines       42236    42270   +34     
=======================================
+ Hits        42113    42147   +34     
  Misses        123      123           

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

[obliviateandsurrender]

Leaving the first set of comments. More to follow later today.

[obliviateandsurrender]

Not a blocker but maybe we can just unify these two to always reshape lol.

[austingmhuang]

You are allowed to mix numpy with any of the other interfaces. So using numpy + jax is considered "the same interface" but using pnp + jax is considered heresy.

[obliviateandsurrender]

Shouldn't these be JAX-arrays?

[austingmhuang]

Oh, if you didn't set the alphas as JAX arrays I purposefully don't coerce them into numpy arrays. numpy arrays are perfectly compatible with jax.