[BUG] Various demos raising jax deprecation warnings

[albi3ro]

Various demos are raising jax deprecation warnings. While they continue to run for now, these demos will fail in the future when jax completes their deprecation cycles. We need to head off these deprecation warnings and fix them.

The demos that raise deprecations include, but are not limited to:

https://pennylane.ai/qml/demos/tutorial_implicit_diff_susceptibility/#implicit-differentiation-of-ground-states-in-pennylane

https://pennylane.ai/qml/demos/tutorial_jax_transformations/

Which both warn:

/home/runner/work/qml/qml/demonstrations/tutorial_implicit_diff_susceptibility.py:266: DeprecationWarning: Accessing jax.config via the jax.config submodule is deprecated.
  from jax.config import config

And:

https://pennylane.ai/qml/demos/tutorial_optimal_control/

which warns

/home/runner/work/qml/qml/demonstrations/tutorial_optimal_control.py:506: DeprecationWarning: hstack requires ndarray or scalar arguments, got <class 'list'> at position 0. In a future JAX release this will be an error.
  params = [jnp.hstack([[0.1 * (-1) ** i for i in range(P)], time]) for time in times]

These are just the demos I manually checked. Other demos may also be raising deprecation warnings uncaught.

[KetpuntoG]

I think it has been solved by @mudit2812

[mudit2812]

@KetpuntoG The warnings I fixed were different. These might still be present.