Return a single ragged list from f_step in JAX mode

[jessegrabowski]

simulate was failing when mode="JAX" because the f_step function was returning two ragged lists instead of just one.

Also adds tests for simulate in JAX mode

---

📚 Documentation preview 📚: https://cge-modeling--22.org.readthedocs.build/en/22/

[codecov]

Codecov Report

Attention: Patch coverage is 73.33333% with 4 lines in your changes missing coverage. Please review.

Project coverage is 68.94%. Comparing base (d6309ef) to head (fa856b6).

Files	Patch %	Lines
cge_modeling/__init__.py	42.85%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #22      +/-   ##
==========================================
+ Coverage   66.30%   68.94%   +2.63%     
==========================================
  Files          17       17              
  Lines        2698     2708      +10     
==========================================
+ Hits         1789     1867      +78     
+ Misses        909      841      -68     

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