test: add test of bool eq api

Author: beckermr

jax_galsim/noise.py

@@ -106,11 +106,15 @@ def _applyTo(self, image):
         raise NotImplementedError("Cannot call applyTo on a pure BaseNoise object")
 
     def __eq__(self, other):
-        # Quick and dirty.  Just check reprs are equal.
-        return self is other or repr(self) == repr(other)
+        if self is other:
+            return jnp.array(True)
+        elif isinstance(other, BaseNoise):
+            return jnp.array(self._rng == other._rng)
+        else:
+            return jnp.array(False)
 
     def __ne__(self, other):
-        return not self.__eq__(other)
+        return ~self.__eq__(other)
 
     __hash__ = None
 
@@ -173,6 +177,16 @@ def __repr__(self):
     def __str__(self):
         return "galsim.GaussianNoise(sigma=%s)" % (ensure_hashable(self.sigma),)
 
+    def __eq__(self, other):
+        if self is other:
+            return jnp.array(True)
+        elif isinstance(other, self.__class__):
+            return jnp.array(self._rng == other._rng) & jnp.array_equal(
+                self._sigma, other._sigma
+            )
+        else:
+            return jnp.array(False)
+
     def tree_flatten(self):
         """This function flattens the GaussianNoise into a list of children
         nodes that will be traced by JAX and auxiliary static data."""
@@ -265,6 +279,16 @@ def __repr__(self):
     def __str__(self):
         return "galsim.PoissonNoise(sky_level=%s)" % (self.sky_level)
 
+    def __eq__(self, other):
+        if self is other:
+            return jnp.array(True)
+        elif isinstance(other, self.__class__):
+            return jnp.array(self._rng == other._rng) & jnp.array_equal(
+                self._sky_level, other._sky_level
+            )
+        else:
+            return jnp.array(False)
+
     def tree_flatten(self):
         """This function flattens the PoissonNoise into a list of children
         nodes that will be traced by JAX and auxiliary static data."""
@@ -429,6 +453,19 @@ def __str__(self):
             self.read_noise,
         )
 
+    def __eq__(self, other):
+        if self is other:
+            return jnp.array(True)
+        elif isinstance(other, self.__class__):
+            return (
+                jnp.array(self._rng == other._rng)
+                & jnp.array_equal(self._sky_level, other._sky_level)
+                & jnp.array_equal(self._gain, other._gain)
+                & jnp.array_equal(self._read_noise, other._read_noise)
+            )
+        else:
+            return jnp.array(False)
+
     def tree_flatten(self):
         """This function flattens the CCDNoise into a list of children
         nodes that will be traced by JAX and auxiliary static data."""
@@ -570,6 +607,16 @@ def __repr__(self):
     def __str__(self):
         return "galsim.VariableGaussianNoise(var_image=%s)" % (self.var_image)
 
+    def __eq__(self, other):
+        if self is other:
+            return jnp.array(True)
+        elif isinstance(other, self.__class__):
+            return jnp.array(self._rng == other._rng) & jnp.array(
+                self._var_image == other._var_image
+            )
+        else:
+            return jnp.array(False)
+
     def tree_flatten(self):
         """This function flattens the VariableGaussianNoise into a list of children
         nodes that will be traced by JAX and auxiliary static data."""

tests/jax/test_api.py

@@ -126,6 +126,16 @@ def _run_object_checks(obj, cls, kind):
 
         # check that we can hash the object
         hash(obj)
+
+        # check that val jax array
+        if (
+            hasattr(obj, "isStatic")
+            and obj.isStatic()
+            or isinstance(obj, jax_galsim.Sensor)
+        ):
+            assert isinstance(eval(repr(obj)) == obj, bool)
+        else:
+            assert isinstance(eval(repr(obj)) == obj, jnp.ndarray)
     elif kind == "to-from-galsim":
         gs_obj = obj.to_galsim()
         jgs_obj = obj.from_galsim(gs_obj)
@@ -141,6 +151,14 @@ def _run_object_checks(obj, cls, kind):
 
         # check that we cannot hash the object
         assert obj.__hash__ is None
+
+        # check that val jax array
+        if (hasattr(obj, "isStatic") and obj.isStatic()) or isinstance(
+            obj, jax_galsim.Sensor
+        ):
+            assert isinstance(eval(repr(obj)) == obj, bool)
+        else:
+            assert isinstance(eval(repr(obj)) == obj, jnp.ndarray)
     elif kind == "pickle-eval-repr-wcs":
         import jax_galsim as galsim  # noqa: F401
 
@@ -152,6 +170,16 @@ def _run_object_checks(obj, cls, kind):
 
         # check that we cannot hash the object
         hash(obj)
+
+        # check that val jax array
+        if (
+            hasattr(obj, "isStatic")
+            and obj.isStatic()
+            or isinstance(obj, jax_galsim.Sensor)
+        ):
+            assert isinstance(eval(repr(obj)) == obj, bool)
+        else:
+            assert isinstance(eval(repr(obj)) == obj, jnp.ndarray)
     elif kind == "jax-compatible":
         # JAX tracing should be an identity
         assert cls.tree_unflatten(*((obj.tree_flatten())[::-1])) == obj