fix: ensure we use any or all in calls for errors

Author: beckermr

jax_galsim/core/utils.py

@@ -45,11 +45,12 @@ def check_is_int_then_cast(val, msg):
         if val != int(val):
             raise TypeError(msg)
         val = int(val)
-    elif not has_tracers(val):
+    else:
         # otherwise we use more opaque checking upon jit via equinox
+        val = jnp.array(val)
         val = equinox.error_if(
             val,
-            val != jnp.trunc(val),
+            np.any(val != jnp.trunc(val)),
             msg,
         )
         val = val.astype(int)

jax_galsim/moffat.py

@@ -74,9 +74,10 @@ def __init__(
                     f"JAX-GalSim does not support Moffat beta values <= {self._beta_thr}."
                 )
         elif not has_tracers(beta):
+            beta = jnp.array(beta)
             beta = equinox.error_if(
-                jnp.array(beta),
-                beta <= self._beta_thr,
+                beta,
+                jnp.any(beta <= self._beta_thr),
                 f"JAX-GalSim does not support Moffat beta values <= {self._beta_thr}.",
             )
 

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 has_tracers, implements
+from jax_galsim.core.utils import implements
 
 try:
     from jax.extend.random import wrap_key_data
@@ -313,10 +313,11 @@ def __init__(self, seed=None, mean=0.0, sigma=1.0):
                 raise ValueError(
                     f"Gaussian deviates must have a non-negative sigma. Got {sigma!r}."
                 )
-        elif not has_tracers(sigma):
+        else:
+            sigma = jnp.array(sigma)
             sigma = equinox.error_if(
-                jnp.array(sigma),
-                sigma < 0,
+                sigma,
+                jnp.any(sigma < 0),
                 f"Gaussian deviates must have a non-negative sigma. Got {sigma!r}.",
             )
 
@@ -466,10 +467,11 @@ def __init__(self, seed=None, mean=1.0):
                 raise ValueError(
                     f"Poisson deviates must have a non-negative mean. Got {mean!r}."
                 )
-        elif not has_tracers(mean):
+        else:
+            mean = jnp.array(mean)
             mean = equinox.error_if(
-                jnp.array(mean),
-                mean < 0,
+                mean,
+                jnp.any(mean < 0),
                 f"Poisson deviates must have a non-negative mean. Got {mean!r}.",
             )