Bugfixes (Python version and Python tests)

[MarcBerneman]

Updated Python versions because Haiku does not support Python 3.8 anymore (google-deepmind/dm-haiku@26de2c0).

FLAGS = jax_config.FLAGS gives an error (see action):
AttributeError: ''Config' object has no attribute 'FLAGS'
This has been fixed by replacing FLAGS.jax_enable_x64 by jax_config.jax_enable_x64

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (e107c26) 82.78% compared to head (e0f58a5) 82.78%.

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

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #295   +/-   ##
=======================================
  Coverage   82.78%   82.78%           
=======================================
  Files          27       27           
  Lines        5245     5245           
=======================================
  Hits         4342     4342           
  Misses        903      903           

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