Review tests

[alicjapolanska]

Tests are getting a bit long to run on actions so profiling and reducing run time would be useful.

Also add jax double precision to make Gaussian-based tests more robust.