fix: more tests

Author: beckermr

tests/jax/test_random_jax.py

@@ -13,10 +13,18 @@ def _make_gauss(sigma):
         _make_gauss(-1.0)
 
     @jax.jit
-    def _make_gauss(sigma):
+    def _make_gauss_again(sigma):
+        return jax_galsim.GaussianDeviate(seed=10, sigma=sigma)
+
+    _make_gauss_again(1.0)
+
+    with pytest.raises(Exception):
+        _make_gauss_again(-1)
+
+    def _make_gauss_again_again(sigma):
         return jax_galsim.GaussianDeviate(seed=10, sigma=sigma)
 
-    _make_gauss(1.0)
+    _make_gauss_again_again(1.0)
 
     with pytest.raises(Exception):
-        _make_gauss(-1)
+        jax.jit(_make_gauss_again_again)(-1)