do the same in jax as in jax-jit for getting results in execute_fwd

[timmysilv]

Context:
The jax code was wrongly iterating over [[res1, res2, ...], jacs] instead of [res1, res2, ...] but it slipped by undetected because you'd need 3 or more tapes to be executed to hit an IndexError.

Description of the Change:
Remove the unnecessary tracing stuff and match what jax-jit does.

Benefits:
Code now works as expected

Possible Drawbacks:
It's a bit hacky, but it gets the job done.

[rmoyard]

Thanks, it looks good to me 👍

[albi3ro]

Also, still need to fix the

TypeError: Custom JVP rule must produce primal and tangent outputs with equal shapes and dtypes, but got:
  primal float32[] for tangent float32[1,0]

that occurs whenever we take a derivative.

[codecov]

Codecov Report

Merging #4190 (4bf86d5) into master (7aa9ce3) will decrease coverage by 0.01%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master    #4190      +/-   ##
==========================================
- Coverage   99.77%   99.77%   -0.01%     
==========================================
  Files         342      342              
  Lines       30704    30696       -8     
==========================================
- Hits        30634    30626       -8     
  Misses         70       70              

Impacted Files	Coverage Δ
pennylane/interfaces/jax.py	99.51% <100.00%> (-0.02%)	⬇️
pennylane/interfaces/jax_jit_tuple.py	100.00% <100.00%> (ø)

[albi3ro]

Thanks for this quick fix.

We can keep trying to track down the a fix for the differentiation, but that can be a later fix.