Jax fixes

Merged Philipp Arras requested to merge jax_fixes into NIFTy_8

Merge request reports