From c373ad608fb47c96b999b018bba01235f0965d15 Mon Sep 17 00:00:00 2001
From: Hossein Pourbozorg <prbzrg@gmail.com>
Date: Tue, 4 Aug 2026 17:40:46 +0330
Subject: [PATCH] update JET and use functions instead of macro

---
 test/Project.toml                       |  2 +-
 test/quality_tests/checkby_JET_tests.jl | 42 ++++++++++---------------
 2 files changed, 17 insertions(+), 27 deletions(-)

diff --git a/test/Project.toml b/test/Project.toml
index 833016b8..c6539325 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -28,7 +28,7 @@ Distributions = "0.25"
 Enzyme = "0.13"
 ExplicitImports = "1"
 ForwardDiff = "1"
-JET = "0.9, 0.10, 0.11"
+JET = "0.9, 0.10, 0.11, 0.12"
 Lux = "1"
 LuxCore = "1"
 MLDataDevices = "1"
diff --git a/test/quality_tests/checkby_JET_tests.jl b/test/quality_tests/checkby_JET_tests.jl
index 1090d790..61d7fbe8 100644
--- a/test/quality_tests/checkby_JET_tests.jl
+++ b/test/quality_tests/checkby_JET_tests.jl
@@ -69,37 +69,27 @@ Test.@testset verbose = true showtiming = true failfast = false "CheckByJET" beg
 
         if conditioned
             ContinuousNormalizingFlows.loss(icnf, omode, r, r2, ps, st)
-            JET.@test_call target_modules = (ContinuousNormalizingFlows,) ContinuousNormalizingFlows.loss(
-                icnf,
-                omode,
-                r,
-                r2,
-                ps,
-                st,
+            JET.test_call(
+                ContinuousNormalizingFlows.loss,
+                typeof((icnf, omode, r, r2, ps, st));
+                target_modules = (ContinuousNormalizingFlows,),
             )
-            JET.@test_opt target_modules = (ContinuousNormalizingFlows,) ContinuousNormalizingFlows.loss(
-                icnf,
-                omode,
-                r,
-                r2,
-                ps,
-                st,
+            JET.test_opt(
+                ContinuousNormalizingFlows.loss,
+                typeof((icnf, omode, r, r2, ps, st));
+                target_modules = (ContinuousNormalizingFlows,),
             )
         else
             ContinuousNormalizingFlows.loss(icnf, omode, r, ps, st)
-            JET.@test_call target_modules = (ContinuousNormalizingFlows,) ContinuousNormalizingFlows.loss(
-                icnf,
-                omode,
-                r,
-                ps,
-                st,
+            JET.test_call(
+                ContinuousNormalizingFlows.loss,
+                typeof((icnf, omode, r, ps, st));
+                target_modules = (ContinuousNormalizingFlows,),
             )
-            JET.@test_opt target_modules = (ContinuousNormalizingFlows,) ContinuousNormalizingFlows.loss(
-                icnf,
-                omode,
-                r,
-                ps,
-                st,
+            JET.test_opt(
+                ContinuousNormalizingFlows.loss,
+                typeof((icnf, omode, r, ps, st));
+                target_modules = (ContinuousNormalizingFlows,),
             )
         end
     end