diff --git a/tests/test_shear_trans.py b/tests/test_shear_trans.py
index 99b3ee0..70cb9b5 100644
--- a/tests/test_shear_trans.py
+++ b/tests/test_shear_trans.py
@@ -7,9 +7,9 @@
 from jax import jit, random
 from jax_galsim import GSParams
 
-from bpd.prior import (
+from bpd.sample import sample_ellip_prior
+from bpd.shear import (
     inv_shear_transformation,
-    sample_ellip_prior,
     scalar_inv_shear_transformation,
     scalar_shear_transformation,
     shear_transformation,