From 8886023fdfbfb0d8b4d9329f8198ed89216d48d2 Mon Sep 17 00:00:00 2001
From: ismael2395 <imendoza@umich.edu>
Date: Mon, 4 Nov 2024 12:17:25 -0800
Subject: [PATCH] need more arguments now

---
 tests/test_shear_trans.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/test_shear_trans.py b/tests/test_shear_trans.py
index 93726f3..da60233 100644
--- 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: