From 95dee3488d9baaba5d449f99e5a04206dc13e500 Mon Sep 17 00:00:00 2001
From: ismael-mendoza <11745764+ismael-mendoza@users.noreply.github.com>
Date: Fri, 18 Sep 2026 12:18:24 -0400
Subject: [PATCH 1/2] use scatter add in test of drawing full scnes for future
 reference

---
 tests/jax/test_render_scene.py | 25 +++++++++++--------------
 1 file changed, 11 insertions(+), 14 deletions(-)

diff --git a/tests/jax/test_render_scene.py b/tests/jax/test_render_scene.py
index 064c9aef..72f90b0d 100644
--- a/tests/jax/test_render_scene.py
+++ b/tests/jax/test_render_scene.py
@@ -183,14 +183,16 @@ def _draw_stamp_jgs(
     return stamp
 
 
-@partial(jax.jit, static_argnames=("slen",))
-def _add_to_image(carry, x, slen):
-    image = carry[0]
-    stamp = x
-
-    image[stamp.bounds] += stamp
-
-    return (image,), None
+def _scatter_stamps_into_image(image: jgs.ImageD, stamps: jgs.ImageD) -> jgs.ImageD:
+    """Add a batch of same-size ``stamps`` into ``image`` with a single vectorized scatter-add."""
+    n, slen, _ = stamps.array.shape
+    row0 = stamps.bounds.ymin - image.bounds.ymin
+    col0 = stamps.bounds.xmin - image.bounds.xmin
+    local = jnp.arange(slen)
+    rows = jnp.broadcast_to(row0[:, None, None] + local[None, :, None], (n, slen, slen))
+    cols = jnp.broadcast_to(col0[:, None, None] + local[None, None, :], (n, slen, slen))
+    new_array = image.array.at[rows, cols].add(stamps.array)
+    return jgs.ImageD(new_array, wcs=image.wcs, bounds=image.bounds)
 
 
 @partial(jax.jit, static_argnames=("fft_size", "slen", "ilen", "ng"))
@@ -217,12 +219,7 @@ def _render_scene_stamps_jax_galsim(
         jnp.pad(image.array, slen), wcs=image.wcs, bounds=image.bounds.withBorder(slen)
     )
 
-    final_pad_image = jax.lax.scan(
-        partial(_add_to_image, slen=slen),
-        (pad_image,),
-        xs=stamps,
-        length=ng,
-    )[0][0]
+    final_pad_image = _scatter_stamps_into_image(pad_image, stamps)
 
     return stamps, final_pad_image
 

From 764e5f23d6f6756c84bd114ce66671f93b459114 Mon Sep 17 00:00:00 2001
From: ismael-mendoza <11745764+ismael-mendoza@users.noreply.github.com>
Date: Fri, 18 Sep 2026 12:21:29 -0400
Subject: [PATCH 2/2] jit like the old one just in case

---
 tests/jax/test_render_scene.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/jax/test_render_scene.py b/tests/jax/test_render_scene.py
index 72f90b0d..580c9d58 100644
--- a/tests/jax/test_render_scene.py
+++ b/tests/jax/test_render_scene.py
@@ -183,6 +183,7 @@ def _draw_stamp_jgs(
     return stamp
 
 
+@jax.jit
 def _scatter_stamps_into_image(image: jgs.ImageD, stamps: jgs.ImageD) -> jgs.ImageD:
     """Add a batch of same-size ``stamps`` into ``image`` with a single vectorized scatter-add."""
     n, slen, _ = stamps.array.shape