fix jax deprecations

[FFroehlich]

• Replaces experimental host_callback calls in jax, which are now deprecated, with pure_callback calls.
• Replaces caching with appropriate calls to the inner objective as caching results in technically non-pure. (Caching the way it was implemented might work anyways, but that's a bit iffy). This required removal of second order code, but I am doubtful that part was functional anyways.
• Removes superfluous calls to jax.jit in Jax.Objective.
• Added compatibility with jax.vmap. This required quite a bit of refactoring as the base objective assumes inputs/outputs to be numpy arrays, which is incompatible with jax batch tracers.
• Refactored tests to actually test interoperability with jax.grad and jax.vmap.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 75.86207% with 7 lines in your changes are missing coverage. Please review.

Project coverage is 84.43%. Comparing base (38d91c4) to head (67ec64f).

Files	Patch %	Lines
pypesto/objective/jax/base.py	75.86%	7 Missing ⚠️

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #1346      +/-   ##
===========================================
- Coverage    84.50%   84.43%   -0.07%     
===========================================
  Files          157      157              
  Lines        12906    12881      -25     
===========================================
- Hits         10906    10876      -30     
- Misses        2000     2005       +5     

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

[PaulJonasJost]

LGTM, not clear on the very details of the implementation though, but the tests checks out 👌🏼 thanks for the update

[FFroehlich]

Now no longer requires passing of input function. Extended tests to demonstrate jax transformations of inputs and outputs.