From 00571c3871542732063f42a58ed5a4df3c9f390f Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Sat, 7 Feb 2026 11:04:43 +0100 Subject: [PATCH 1/3] minor edits for updated comp --- jax_galsim/fitswcs.py | 9 +++------ jax_galsim/integ.py | 2 +- pyproject.toml | 5 ++--- 3 files changed, 6 insertions(+), 10 deletions(-) diff --git a/jax_galsim/fitswcs.py b/jax_galsim/fitswcs.py index 5870d755..c4756f74 100644 --- a/jax_galsim/fitswcs.py +++ b/jax_galsim/fitswcs.py @@ -833,17 +833,14 @@ def __eq__(self, other): and jnp.array_equal(self.crpix, other.crpix) and jnp.array_equal(self.cd, other.cd) and self.center == other.center - and ( + and ((self.pv is None and other.pv is None) or jnp.array_equal(self.pv, other.pv) - or (self.pv is None and other.pv is None) ) - and ( + and ( (self.ab is None and other.ab is None) or jnp.array_equal(self.ab, other.ab) - or (self.ab is None and other.ab is None) ) - and ( + and ((self.abp is None and other.abp is None) or jnp.array_equal(self.abp, other.abp) - or (self.abp is None and other.abp is None) ) ) diff --git a/jax_galsim/integ.py b/jax_galsim/integ.py index c2e19a42..19ad5c4b 100644 --- a/jax_galsim/integ.py +++ b/jax_galsim/integ.py @@ -40,7 +40,7 @@ def int1d( @jax.jit def _func(x): rdt = jax.ShapeDtypeStruct(x.shape, x.dtype) - return jax.pure_callback(func, rdt, x) + return jax.pure_callback(func, rdt, x, vmap_method="sequential") else: _func = func diff --git a/pyproject.toml b/pyproject.toml index ff8a2390..9a44d112 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,10 +16,9 @@ readme = "README.md" dependencies = [ "numpy >=1.18.0", "galsim >=2.3.0", - "jax <0.7.0", - "jaxlib", + "jax >=0.4.0", "astropy >=2.0", - "tensorflow-probability >=0.21.0", + "tfp-nightly", "quadax", ] From 137dda9b252635885d29f870b9bca78036ed5233 Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Sat, 7 Feb 2026 11:11:42 +0100 Subject: [PATCH 2/3] fixing formatting --- jax_galsim/fitswcs.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/jax_galsim/fitswcs.py b/jax_galsim/fitswcs.py index c4756f74..232801d6 100644 --- a/jax_galsim/fitswcs.py +++ b/jax_galsim/fitswcs.py @@ -833,14 +833,17 @@ def __eq__(self, other): and jnp.array_equal(self.crpix, other.crpix) and jnp.array_equal(self.cd, other.cd) and self.center == other.center - and ((self.pv is None and other.pv is None) or - jnp.array_equal(self.pv, other.pv) + and ( + (self.pv is None and other.pv is None) + or jnp.array_equal(self.pv, other.pv) ) - and ( (self.ab is None and other.ab is None) or - jnp.array_equal(self.ab, other.ab) + and ( + (self.ab is None and other.ab is None) + or jnp.array_equal(self.ab, other.ab) ) - and ((self.abp is None and other.abp is None) or - jnp.array_equal(self.abp, other.abp) + and ( + (self.abp is None and other.abp is None) + or jnp.array_equal(self.abp, other.abp) ) ) From dca40a7d91e1a31564dd2ad4f859a5f769e63721 Mon Sep 17 00:00:00 2001 From: Francois Lanusse Date: Sat, 7 Feb 2026 12:59:51 +0100 Subject: [PATCH 3/3] Update pyproject.toml Co-authored-by: Matthew R. Becker --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 9a44d112..1921ba52 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ readme = "README.md" dependencies = [ "numpy >=1.18.0", "galsim >=2.3.0", - "jax >=0.4.0", + "jax >=0.8.0", "astropy >=2.0", "tfp-nightly", "quadax",