need more arguments now

LSSTDESC · Nov 4, 2024 · 8886023 · 8886023
1 parent 0f62810
commit 8886023
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/test_shear_trans.py b/tests/test_shear_trans.py
@@ -1,3 +1,5 @@
+from functools import partial
+
 import numpy as np
 import pytest
 from jax import jit as jjit
@@ -53,7 +55,7 @@ def test_image_shear_commute():
     hlr = 0.9
     x, y = (1, 1)
 
-    draw_jitted = jjit(draw_gaussian)
+    draw_jitted = jjit(partial(draw_gaussian), slen=53, fft_size=256)
     for e1 in ellips:
         for e2 in ellips:
             for g1 in shears: