fix: apparently this does not work on tracers

Author: beckermr

jax_galsim/random.py

@@ -9,7 +9,7 @@
 import numpy as np
 from jax.tree_util import register_pytree_node_class
 
-from jax_galsim.core.utils import implements
+from jax_galsim.core.utils import has_tracers, implements
 
 try:
     from jax.extend.random import wrap_key_data
@@ -313,7 +313,7 @@ def __init__(self, seed=None, mean=0.0, sigma=1.0):
                 raise ValueError(
                     f"Gaussian deviates must have a positive sigma. Got {sigma!r}."
                 )
-        else:
+        elif not has_tracers(sigma):
             sigma = equinox.error_if(
                 jnp.array(sigma),
                 sigma <= 0,
@@ -466,7 +466,7 @@ def __init__(self, seed=None, mean=1.0):
                 raise ValueError(
                     f"Poisson deviates must have a non-negative mean. Got {mean!r}."
                 )
-        else:
+        elif not has_tracers(mean):
             mean = equinox.error_if(
                 jnp.array(mean),
                 mean < 0,

tests/jax/test_random_jax.py

@@ -0,0 +1,22 @@
+import jax
+import pytest
+
+import jax_galsim
+
+
+def test_random_jax_gaussian_pos_sigma_jit():
+    @jax.jit
+    def _make_gauss(sigma):
+        return jax_galsim.GaussianDeviate(seed=10, sigma=sigma)
+
+    with pytest.raises(Exception):
+        _make_gauss(-1.0)
+
+    @jax.jit
+    def _make_gauss(sigma):
+        return jax_galsim.GaussianDeviate(seed=10, sigma=sigma)
+
+    _make_gauss(1.0)
+
+    with pytest.raises(Exception):
+        _make_gauss(-1)