From 9d726b1188efafa79b8096cbacc0ab7afa70e08d Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Mon, 6 Mar 2023 11:49:15 +0100
Subject: [PATCH 01/21] (fix): change to `static_argnums` to work around
 decorator

---
 torchquad/integration/newton_cotes.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchquad/integration/newton_cotes.py b/torchquad/integration/newton_cotes.py
index 3e0f2741..c3a4d2af 100644
--- a/torchquad/integration/newton_cotes.py
+++ b/torchquad/integration/newton_cotes.py
@@ -148,7 +148,7 @@ def get_jit_compiled_integrate(
                         self.calculate_grid, static_argnames=["N"]
                     )
                     self._jax_jit_calculate_result = jax.jit(
-                        self.calculate_result, static_argnames=["dim", "n_per_dim"]
+                        self.calculate_result, static_argnums=(1, 2,) # dim and n_per_dim
                     )
                 jit_calculate_grid = self._jax_jit_calculate_grid
                 jit_calculate_result = self._jax_jit_calculate_result

From 93dd533d901a1f36f7a086edaf6581bd6489e74d Mon Sep 17 00:00:00 2001
From: htoftevaag <htoftevaag@users.noreply.github.com>
Date: Mon, 6 Mar 2023 10:53:32 +0000
Subject: [PATCH 02/21] WORKFLOW: Format Python code with black

---
 torchquad/integration/newton_cotes.py | 6 +++++-
 torchquad/tests/boole_test.py         | 3 ++-
 torchquad/utils/set_up_backend.py     | 3 ++-
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/torchquad/integration/newton_cotes.py b/torchquad/integration/newton_cotes.py
index c3a4d2af..9fac80dd 100644
--- a/torchquad/integration/newton_cotes.py
+++ b/torchquad/integration/newton_cotes.py
@@ -148,7 +148,11 @@ def get_jit_compiled_integrate(
                         self.calculate_grid, static_argnames=["N"]
                     )
                     self._jax_jit_calculate_result = jax.jit(
-                        self.calculate_result, static_argnums=(1, 2,) # dim and n_per_dim
+                        self.calculate_result,
+                        static_argnums=(
+                            1,
+                            2,
+                        ),  # dim and n_per_dim
                     )
                 jit_calculate_grid = self._jax_jit_calculate_grid
                 jit_calculate_result = self._jax_jit_calculate_result
diff --git a/torchquad/tests/boole_test.py b/torchquad/tests/boole_test.py
index 4e42f3ca..beeeace3 100644
--- a/torchquad/tests/boole_test.py
+++ b/torchquad/tests/boole_test.py
@@ -13,7 +13,8 @@
 
 def _run_boole_tests(backend, _precision):
     """Test the integrate function in integration.Boole for the given backend.
-    Note: For now the 10-D test is diabled due to lack of GPU memory on some computers."""
+    Note: For now the 10-D test is diabled due to lack of GPU memory on some computers.
+    """
 
     bl = Boole()
     # 1D Tests
diff --git a/torchquad/utils/set_up_backend.py b/torchquad/utils/set_up_backend.py
index ec77c9d1..8d88bfff 100644
--- a/torchquad/utils/set_up_backend.py
+++ b/torchquad/utils/set_up_backend.py
@@ -7,7 +7,8 @@
 
 def _get_default_backend():
     """Get the latest backend which was passed to set_up_backend.
-    If set_up_backend has never been executed, return "torch" for backwards compatibility"""
+    If set_up_backend has never been executed, return "torch" for backwards compatibility
+    """
     return os.environ.get("TORCHQUAD_DEFAULT_BACKEND", "torch")
 
 

From 050fe01383eda87fcf2aa07ccd9692ce5b74381a Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Mon, 6 Mar 2023 16:06:12 +0100
Subject: [PATCH 03/21] (feat): add newton-cotest integrator tests

---
 torchquad/tests/boole_test.py     | 21 +++++++++++++++++++++
 torchquad/tests/simpson_test.py   | 20 ++++++++++++++++++++
 torchquad/tests/trapezoid_test.py | 20 ++++++++++++++++++++
 3 files changed, 61 insertions(+)

diff --git a/torchquad/tests/boole_test.py b/torchquad/tests/boole_test.py
index beeeace3..810ae78b 100644
--- a/torchquad/tests/boole_test.py
+++ b/torchquad/tests/boole_test.py
@@ -61,6 +61,27 @@ def _run_boole_tests(backend, _precision):
     # for error in errors:
     # assert error < 5e-9
 
+    # JIT Tests
+    if backend != "numpy":
+        N = 401
+
+        integrate = bl.get_jit_compiled_integrate(
+            dim=1, N=N, backend=backend
+        )
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+        )
+        print(f"1D Boole JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+        # Polynomials up to degree 5 can be integrated almost exactly with Boole.
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 5 or err < 6.33e-11
+        for error in errors:
+            assert error < 6.33e-11
+
 
 test_integrate_numpy = setup_test_for_backend(_run_boole_tests, "numpy", "float64")
 test_integrate_torch = setup_test_for_backend(_run_boole_tests, "torch", "float64")
diff --git a/torchquad/tests/simpson_test.py b/torchquad/tests/simpson_test.py
index 28172c56..23e3afe0 100644
--- a/torchquad/tests/simpson_test.py
+++ b/torchquad/tests/simpson_test.py
@@ -89,6 +89,26 @@ def _run_simpson_tests(backend, _precision):
     for error in errors:
         assert error < 5e-9
 
+    if backend != "numpy":
+        N = 401
+
+        integrate = simp.get_jit_compiled_integrate(
+            dim=1, N=N, backend=backend
+        )
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+        )
+        print(f"1D Simpson JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+        # Polynomials up to degree 5 can be integrated almost exactly with Boole.
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 5 or err < 6.33e-11
+        for error in errors:
+            assert error < 6.33e-11
+
 
 test_integrate_numpy = setup_test_for_backend(_run_simpson_tests, "numpy", "float64")
 test_integrate_torch = setup_test_for_backend(_run_simpson_tests, "torch", "float64")
diff --git a/torchquad/tests/trapezoid_test.py b/torchquad/tests/trapezoid_test.py
index 0d224d1a..c863a5fc 100644
--- a/torchquad/tests/trapezoid_test.py
+++ b/torchquad/tests/trapezoid_test.py
@@ -86,6 +86,26 @@ def _run_trapezoid_tests(backend, _precision):
     for error in errors:
         assert error < 7000
 
+    if backend != "numpy":
+        N = 401
+
+        integrate = tp.get_jit_compiled_integrate(
+            dim=1, N=N, backend=backend
+        )
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+        )
+        print(f"1D Trapezoid JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+        # Polynomials up to degree 5 can be integrated almost exactly with Boole.
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 5 or err < 6.33e-11
+        for error in errors:
+            assert error < 6.33e-11
+
 
 test_integrate_numpy = setup_test_for_backend(_run_trapezoid_tests, "numpy", "float64")
 test_integrate_torch = setup_test_for_backend(_run_trapezoid_tests, "torch", "float64")

From 41c6f9d1f4429a3b82c0d7a22a6b17bf8165a894 Mon Sep 17 00:00:00 2001
From: htoftevaag <htoftevaag@users.noreply.github.com>
Date: Mon, 6 Mar 2023 15:07:44 +0000
Subject: [PATCH 04/21] WORKFLOW: Format Python code with black

---
 torchquad/tests/boole_test.py     | 4 +---
 torchquad/tests/simpson_test.py   | 8 ++++----
 torchquad/tests/trapezoid_test.py | 8 ++++----
 3 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/torchquad/tests/boole_test.py b/torchquad/tests/boole_test.py
index 810ae78b..72350ac8 100644
--- a/torchquad/tests/boole_test.py
+++ b/torchquad/tests/boole_test.py
@@ -65,9 +65,7 @@ def _run_boole_tests(backend, _precision):
     if backend != "numpy":
         N = 401
 
-        integrate = bl.get_jit_compiled_integrate(
-            dim=1, N=N, backend=backend
-        )
+        integrate = bl.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
diff --git a/torchquad/tests/simpson_test.py b/torchquad/tests/simpson_test.py
index 23e3afe0..c6dcf4ca 100644
--- a/torchquad/tests/simpson_test.py
+++ b/torchquad/tests/simpson_test.py
@@ -92,9 +92,7 @@ def _run_simpson_tests(backend, _precision):
     if backend != "numpy":
         N = 401
 
-        integrate = simp.get_jit_compiled_integrate(
-            dim=1, N=N, backend=backend
-        )
+        integrate = simp.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
@@ -102,7 +100,9 @@ def _run_simpson_tests(backend, _precision):
             use_complex=True,
             backend=backend,
         )
-        print(f"1D Simpson JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+        print(
+            f"1D Simpson JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}"
+        )
         # Polynomials up to degree 5 can be integrated almost exactly with Boole.
         for err, test_function in zip(errors, funcs):
             assert test_function.get_order() > 5 or err < 6.33e-11
diff --git a/torchquad/tests/trapezoid_test.py b/torchquad/tests/trapezoid_test.py
index c863a5fc..466e3a92 100644
--- a/torchquad/tests/trapezoid_test.py
+++ b/torchquad/tests/trapezoid_test.py
@@ -89,9 +89,7 @@ def _run_trapezoid_tests(backend, _precision):
     if backend != "numpy":
         N = 401
 
-        integrate = tp.get_jit_compiled_integrate(
-            dim=1, N=N, backend=backend
-        )
+        integrate = tp.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
@@ -99,7 +97,9 @@ def _run_trapezoid_tests(backend, _precision):
             use_complex=True,
             backend=backend,
         )
-        print(f"1D Trapezoid JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+        print(
+            f"1D Trapezoid JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}"
+        )
         # Polynomials up to degree 5 can be integrated almost exactly with Boole.
         for err, test_function in zip(errors, funcs):
             assert test_function.get_order() > 5 or err < 6.33e-11

From 694b724e28c2d7b8014fbbd70d684a61c5a048d7 Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Mon, 6 Mar 2023 16:26:29 +0100
Subject: [PATCH 05/21] (fix): re-compile per integrand because of different
 shapes

---
 torchquad/tests/boole_test.py     |  9 +++++----
 torchquad/tests/simpson_test.py   | 16 ++++++++++------
 torchquad/tests/trapezoid_test.py | 16 ++++++++++------
 3 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/torchquad/tests/boole_test.py b/torchquad/tests/boole_test.py
index 810ae78b..35dae818 100644
--- a/torchquad/tests/boole_test.py
+++ b/torchquad/tests/boole_test.py
@@ -64,10 +64,11 @@ def _run_boole_tests(backend, _precision):
     # JIT Tests
     if backend != "numpy":
         N = 401
-
-        integrate = bl.get_jit_compiled_integrate(
-            dim=1, N=N, backend=backend
-        )
+        def integrate(*args, **kwargs):
+            jit_integrate = bl.get_jit_compiled_integrate(
+                dim=1, N=N, backend=backend
+            )
+            return jit_integrate(*args, **kwargs)
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
diff --git a/torchquad/tests/simpson_test.py b/torchquad/tests/simpson_test.py
index 23e3afe0..e5208ff7 100644
--- a/torchquad/tests/simpson_test.py
+++ b/torchquad/tests/simpson_test.py
@@ -90,11 +90,13 @@ def _run_simpson_tests(backend, _precision):
         assert error < 5e-9
 
     if backend != "numpy":
-        N = 401
+        N = 100001
 
-        integrate = simp.get_jit_compiled_integrate(
-            dim=1, N=N, backend=backend
-        )
+        def integrate(*args, **kwargs):
+            jit_integrate = simp.get_jit_compiled_integrate(
+                dim=1, N=N, backend=backend
+            )
+            return jit_integrate(*args, **kwargs)
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
@@ -105,9 +107,11 @@ def _run_simpson_tests(backend, _precision):
         print(f"1D Simpson JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
         # Polynomials up to degree 5 can be integrated almost exactly with Boole.
         for err, test_function in zip(errors, funcs):
-            assert test_function.get_order() > 5 or err < 6.33e-11
+            assert test_function.get_order() > 3 or (
+                err < 3e-11 if test_function.is_integrand_1d else err < 6e-10
+            )  # errors add up if the integrand is higher dimensional
         for error in errors:
-            assert error < 6.33e-11
+            assert error < 1e-7
 
 
 test_integrate_numpy = setup_test_for_backend(_run_simpson_tests, "numpy", "float64")
diff --git a/torchquad/tests/trapezoid_test.py b/torchquad/tests/trapezoid_test.py
index c863a5fc..f466f651 100644
--- a/torchquad/tests/trapezoid_test.py
+++ b/torchquad/tests/trapezoid_test.py
@@ -87,11 +87,13 @@ def _run_trapezoid_tests(backend, _precision):
         assert error < 7000
 
     if backend != "numpy":
-        N = 401
+        N = 100000
 
-        integrate = tp.get_jit_compiled_integrate(
-            dim=1, N=N, backend=backend
-        )
+        def integrate(*args, **kwargs):
+            jit_integrate = tp.get_jit_compiled_integrate(
+                dim=1, N=N, backend=backend
+            )
+            return jit_integrate(*args, **kwargs)
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
@@ -102,9 +104,11 @@ def _run_trapezoid_tests(backend, _precision):
         print(f"1D Trapezoid JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
         # Polynomials up to degree 5 can be integrated almost exactly with Boole.
         for err, test_function in zip(errors, funcs):
-            assert test_function.get_order() > 5 or err < 6.33e-11
+            assert test_function.get_order() > 1 or (
+                err < 2e-11 if test_function.is_integrand_1d else err < 5e-10
+            )  # errors add up if the integrand is higher dimensional
         for error in errors:
-            assert error < 6.33e-11
+            assert error < 1e-5
 
 
 test_integrate_numpy = setup_test_for_backend(_run_trapezoid_tests, "numpy", "float64")

From ae99b91bdbc751601d8dd8cc92f1b06e94be6a68 Mon Sep 17 00:00:00 2001
From: htoftevaag <htoftevaag@users.noreply.github.com>
Date: Mon, 6 Mar 2023 15:27:47 +0000
Subject: [PATCH 06/21] WORKFLOW: Format Python code with black

---
 torchquad/tests/boole_test.py     | 6 +++---
 torchquad/tests/simpson_test.py   | 5 ++---
 torchquad/tests/trapezoid_test.py | 5 ++---
 3 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/torchquad/tests/boole_test.py b/torchquad/tests/boole_test.py
index 35dae818..50abb846 100644
--- a/torchquad/tests/boole_test.py
+++ b/torchquad/tests/boole_test.py
@@ -64,11 +64,11 @@ def _run_boole_tests(backend, _precision):
     # JIT Tests
     if backend != "numpy":
         N = 401
+
         def integrate(*args, **kwargs):
-            jit_integrate = bl.get_jit_compiled_integrate(
-                dim=1, N=N, backend=backend
-            )
+            jit_integrate = bl.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
             return jit_integrate(*args, **kwargs)
+
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
diff --git a/torchquad/tests/simpson_test.py b/torchquad/tests/simpson_test.py
index b6410638..398b4c42 100644
--- a/torchquad/tests/simpson_test.py
+++ b/torchquad/tests/simpson_test.py
@@ -93,10 +93,9 @@ def _run_simpson_tests(backend, _precision):
         N = 100001
 
         def integrate(*args, **kwargs):
-            jit_integrate = simp.get_jit_compiled_integrate(
-                dim=1, N=N, backend=backend
-            )
+            jit_integrate = simp.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
             return jit_integrate(*args, **kwargs)
+
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
diff --git a/torchquad/tests/trapezoid_test.py b/torchquad/tests/trapezoid_test.py
index 46a6eb7b..0d869dd1 100644
--- a/torchquad/tests/trapezoid_test.py
+++ b/torchquad/tests/trapezoid_test.py
@@ -90,10 +90,9 @@ def _run_trapezoid_tests(backend, _precision):
         N = 100000
 
         def integrate(*args, **kwargs):
-            jit_integrate = tp.get_jit_compiled_integrate(
-                dim=1, N=N, backend=backend
-            )
+            jit_integrate = tp.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
             return jit_integrate(*args, **kwargs)
+
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},

From 9b60042cc423f8dbd98c7abe37b97272b91e7877 Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Mon, 6 Mar 2023 16:51:50 +0100
Subject: [PATCH 07/21] (chore): remove erroneous boole comments

---
 torchquad/tests/simpson_test.py   | 1 -
 torchquad/tests/trapezoid_test.py | 1 -
 2 files changed, 2 deletions(-)

diff --git a/torchquad/tests/simpson_test.py b/torchquad/tests/simpson_test.py
index 398b4c42..629906fc 100644
--- a/torchquad/tests/simpson_test.py
+++ b/torchquad/tests/simpson_test.py
@@ -106,7 +106,6 @@ def integrate(*args, **kwargs):
         print(
             f"1D Simpson JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}"
         )
-        # Polynomials up to degree 5 can be integrated almost exactly with Boole.
         for err, test_function in zip(errors, funcs):
             assert test_function.get_order() > 3 or (
                 err < 3e-11 if test_function.is_integrand_1d else err < 6e-10
diff --git a/torchquad/tests/trapezoid_test.py b/torchquad/tests/trapezoid_test.py
index 0d869dd1..e9853ac6 100644
--- a/torchquad/tests/trapezoid_test.py
+++ b/torchquad/tests/trapezoid_test.py
@@ -103,7 +103,6 @@ def integrate(*args, **kwargs):
         print(
             f"1D Trapezoid JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}"
         )
-        # Polynomials up to degree 5 can be integrated almost exactly with Boole.
         for err, test_function in zip(errors, funcs):
             assert test_function.get_order() > 1 or (
                 err < 2e-11 if test_function.is_integrand_1d else err < 5e-10

From 1461a1b53029a04ec143efaa6bbc0f06244cd477 Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Mon, 6 Mar 2023 16:55:10 +0100
Subject: [PATCH 08/21] (fix): add mc test

---
 torchquad/tests/monte_carlo_test.py | 33 +++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/torchquad/tests/monte_carlo_test.py b/torchquad/tests/monte_carlo_test.py
index ab624bb0..8788d292 100644
--- a/torchquad/tests/monte_carlo_test.py
+++ b/torchquad/tests/monte_carlo_test.py
@@ -82,6 +82,39 @@ def _run_monte_carlo_tests(backend, _precision):
     for error in errors:
         assert error < 26
 
+    # JIT Tests
+    if backend != "numpy":
+        N = 100000
+
+        def integrate(*args, **kwargs):
+            jit_integrate = mc.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
+            return jit_integrate(*args, **kwargs)
+
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+        )
+        print(f"1D MC JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 0 or err < 1e-14
+
+        # If this breaks check if test functions in helper_functions changed.
+        for error in errors[:3]:
+            assert error < 7e-3
+
+        assert errors[3] < 0.5
+        assert errors[4] < 32.0
+
+        for error in errors[6:10]:
+            assert error < 2e-2
+
+        for error in errors[10:]:
+            assert error < 28.03
+
 
 test_integrate_numpy = setup_test_for_backend(
     _run_monte_carlo_tests, "numpy", "float32"

From ebb2edd4bad6d22186930975a30fc434212e0e65 Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Mon, 6 Mar 2023 17:02:48 +0100
Subject: [PATCH 09/21] (chore): add note about `integrand` shape

---
 docs/source/tutorial.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/source/tutorial.rst b/docs/source/tutorial.rst
index d2164e21..c9d034d1 100644
--- a/docs/source/tutorial.rst
+++ b/docs/source/tutorial.rst
@@ -586,7 +586,7 @@ Compiling the integrate method
 ``````````````````````````````
 
 To speed up the quadrature in situations where it is executed often with the
-same number of points ``N`` and dimensionality ``dim``,
+same number of points ``N``, dimensionality ``dim``, and shape of the ``integrand``,
 we can JIT-compile the performance-relevant parts of the integrate method:
 
 .. code:: ipython3

From 97f067a4202e550d7465d2d98f0810b2a49a5a3a Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Tue, 7 Mar 2023 14:07:00 +0100
Subject: [PATCH 10/21] (feat): switch to re-use fo the jit integral

---
 torchquad/tests/boole_test.py       | 41 +++++++++++++++++++++++++++--
 torchquad/tests/helper_functions.py |  8 +++---
 torchquad/tests/monte_carlo_test.py | 41 ++++++++++++++++++++++++++---
 torchquad/tests/simpson_test.py     | 37 +++++++++++++++++++++++++-
 torchquad/tests/trapezoid_test.py   | 39 +++++++++++++++++++++++++--
 5 files changed, 155 insertions(+), 11 deletions(-)

diff --git a/torchquad/tests/boole_test.py b/torchquad/tests/boole_test.py
index 50abb846..a2025543 100644
--- a/torchquad/tests/boole_test.py
+++ b/torchquad/tests/boole_test.py
@@ -64,9 +64,14 @@ def _run_boole_tests(backend, _precision):
     # JIT Tests
     if backend != "numpy":
         N = 401
+        jit_integrate = None
 
         def integrate(*args, **kwargs):
-            jit_integrate = bl.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = bl.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
             return jit_integrate(*args, **kwargs)
 
         errors, funcs = compute_integration_test_errors(
@@ -75,8 +80,40 @@ def integrate(*args, **kwargs):
             integration_dim=1,
             use_complex=True,
             backend=backend,
+            filter_test_functions=lambda x: x.is_integrand_1d,
+        )
+        print(
+            f"1D Boole JIT Test passed for 1D integrands. N: {N}, backend: {backend}, Errors: {errors}"
+        )
+        # Polynomials up to degree 5 can be integrated almost exactly with Boole.
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 5 or err < 6.33e-11
+        for error in errors:
+            assert error < 6.33e-11
+
+        jit_integrate = (
+            None  # set to None again so can be re-used with new integrand shape
+        )
+
+        def integrate(*args, **kwargs):
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = bl.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
+            return jit_integrate(*args, **kwargs)
+
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+            filter_test_functions=lambda x: x.integrand_dims == [2, 2, 2],
+        )
+        print(
+            f"1D Boole JIT Test passed for [2, 2, 2] dimensional integrands. N: {N}, backend: {backend}, Errors: {errors}"
         )
-        print(f"1D Boole JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
         # Polynomials up to degree 5 can be integrated almost exactly with Boole.
         for err, test_function in zip(errors, funcs):
             assert test_function.get_order() > 5 or err < 6.33e-11
diff --git a/torchquad/tests/helper_functions.py b/torchquad/tests/helper_functions.py
index c0dcc884..88a8ee84 100644
--- a/torchquad/tests/helper_functions.py
+++ b/torchquad/tests/helper_functions.py
@@ -324,6 +324,7 @@ def compute_integration_test_errors(
     use_complex,
     backend,
     use_multi_dim_integrand=True,
+    filter_test_functions=lambda x: x,
 ):
     """Computes errors on all test functions for given dimension and integrator.
 
@@ -334,7 +335,7 @@ def compute_integration_test_errors(
         use_complex (Boolean): If True, skip complex example functions.
         backend (string): Numerical backend for the example functions.
         use_multi_dim_integrand (bool, optional): Whether or not to allow for a multi-dimensional integrand i.e an array of integrands
-
+        filter_test_functions (function, optional): function for filtering which test functions to run
     Returns:
         (list, list): Absolute errors on all example functions and the chosen
             example functions
@@ -344,8 +345,9 @@ def compute_integration_test_errors(
 
     # Compute integration errors on the chosen functions and remember those
     # functions
-    for test_function in get_test_functions(
-        integration_dim, backend, use_multi_dim_integrand
+    for test_function in filter(
+        filter_test_functions,
+        get_test_functions(integration_dim, backend, use_multi_dim_integrand),
     ):
         if not use_complex and test_function.is_complex:
             continue
diff --git a/torchquad/tests/monte_carlo_test.py b/torchquad/tests/monte_carlo_test.py
index 8788d292..209cb550 100644
--- a/torchquad/tests/monte_carlo_test.py
+++ b/torchquad/tests/monte_carlo_test.py
@@ -85,9 +85,14 @@ def _run_monte_carlo_tests(backend, _precision):
     # JIT Tests
     if backend != "numpy":
         N = 100000
+        jit_integrate = None
 
         def integrate(*args, **kwargs):
-            jit_integrate = mc.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = mc.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
             return jit_integrate(*args, **kwargs)
 
         errors, funcs = compute_integration_test_errors(
@@ -96,15 +101,18 @@ def integrate(*args, **kwargs):
             integration_dim=1,
             use_complex=True,
             backend=backend,
+            filter_test_functions=lambda x: x.is_integrand_1d,
+        )
+        print(
+            f"1D MC JIT Test passed for 1D integrands. N: {N}, backend: {backend}, Errors: {errors}"
         )
-        print(f"1D MC JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
 
         for err, test_function in zip(errors, funcs):
             assert test_function.get_order() > 0 or err < 1e-14
 
         # If this breaks check if test functions in helper_functions changed.
         for error in errors[:3]:
-            assert error < 7e-3
+            assert error < 1e-2
 
         assert errors[3] < 0.5
         assert errors[4] < 32.0
@@ -115,6 +123,33 @@ def integrate(*args, **kwargs):
         for error in errors[10:]:
             assert error < 28.03
 
+        jit_integrate = (
+            None  # set to None again so can be re-used with new integrand shape
+        )
+
+        def integrate(*args, **kwargs):
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = mc.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
+            return jit_integrate(*args, **kwargs)
+
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+            filter_test_functions=lambda x: x.integrand_dims == [2, 2, 2],
+        )
+        print(
+            f"1D MC JIT Test passed for [2, 2, 2] dimensional integrands. N: {N}, backend: {backend}, Errors: {errors}"
+        )
+
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 0 or err < 1e-14
+
 
 test_integrate_numpy = setup_test_for_backend(
     _run_monte_carlo_tests, "numpy", "float32"
diff --git a/torchquad/tests/simpson_test.py b/torchquad/tests/simpson_test.py
index 629906fc..c6dbe3c4 100644
--- a/torchquad/tests/simpson_test.py
+++ b/torchquad/tests/simpson_test.py
@@ -91,9 +91,14 @@ def _run_simpson_tests(backend, _precision):
 
     if backend != "numpy":
         N = 100001
+        jit_integrate = None
 
         def integrate(*args, **kwargs):
-            jit_integrate = simp.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = simp.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
             return jit_integrate(*args, **kwargs)
 
         errors, funcs = compute_integration_test_errors(
@@ -102,7 +107,9 @@ def integrate(*args, **kwargs):
             integration_dim=1,
             use_complex=True,
             backend=backend,
+            filter_test_functions=lambda x: x.is_integrand_1d,
         )
+
         print(
             f"1D Simpson JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}"
         )
@@ -113,6 +120,34 @@ def integrate(*args, **kwargs):
         for error in errors:
             assert error < 1e-7
 
+        jit_integrate = None
+
+        def integrate(*args, **kwargs):
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = simp.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
+            return jit_integrate(*args, **kwargs)
+
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+            filter_test_functions=lambda x: x.integrand_dims == [2, 2, 2],
+        )
+        print(
+            f"1D Simpson JIT Test passed for [2, 2, 2] dimensional integrands. N: {N}, backend: {backend}, Errors: {errors}"
+        )
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 3 or (
+                err < 3e-11 if test_function.is_integrand_1d else err < 6e-10
+            )  # errors add up if the integrand is higher dimensional
+        for error in errors:
+            assert error < 1e-7
+
 
 test_integrate_numpy = setup_test_for_backend(_run_simpson_tests, "numpy", "float64")
 test_integrate_torch = setup_test_for_backend(_run_simpson_tests, "torch", "float64")
diff --git a/torchquad/tests/trapezoid_test.py b/torchquad/tests/trapezoid_test.py
index e9853ac6..11116253 100644
--- a/torchquad/tests/trapezoid_test.py
+++ b/torchquad/tests/trapezoid_test.py
@@ -88,9 +88,14 @@ def _run_trapezoid_tests(backend, _precision):
 
     if backend != "numpy":
         N = 100000
+        jit_integrate = None
 
         def integrate(*args, **kwargs):
-            jit_integrate = tp.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = tp.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
             return jit_integrate(*args, **kwargs)
 
         errors, funcs = compute_integration_test_errors(
@@ -99,9 +104,39 @@ def integrate(*args, **kwargs):
             integration_dim=1,
             use_complex=True,
             backend=backend,
+            filter_test_functions=lambda x: x.is_integrand_1d,
+        )
+
+        print(
+            f"1D Trapezoid JIT Test passed for 1D integrands. N: {N}, backend: {backend}, Errors: {errors}"
+        )
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 1 or (
+                err < 2e-11 if test_function.is_integrand_1d else err < 5e-10
+            )  # errors add up if the integrand is higher dimensional
+        for error in errors:
+            assert error < 1e-5
+
+        jit_integrate = None
+
+        def integrate(*args, **kwargs):
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = tp.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
+            return jit_integrate(*args, **kwargs)
+
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+            filter_test_functions=lambda x: x.integrand_dims == [2, 2, 2],
         )
         print(
-            f"1D Trapezoid JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}"
+            f"1D Trapezoid JIT Test passed for [2, 2, 2] dimensional integrands. N: {N}, backend: {backend}, Errors: {errors}"
         )
         for err, test_function in zip(errors, funcs):
             assert test_function.get_order() > 1 or (

From 7b641154c40a8aef356c556d6244042180c54403 Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Wed, 8 Mar 2023 10:54:06 +0100
Subject: [PATCH 11/21] (chore): address comments

---
 torchquad/tests/boole_test.py       | 10 ++--------
 torchquad/tests/monte_carlo_test.py | 10 ++--------
 torchquad/tests/simpson_test.py     | 15 ++++++---------
 torchquad/tests/trapezoid_test.py   | 15 ++++++---------
 4 files changed, 16 insertions(+), 34 deletions(-)

diff --git a/torchquad/tests/boole_test.py b/torchquad/tests/boole_test.py
index a2025543..aaea3f1a 100644
--- a/torchquad/tests/boole_test.py
+++ b/torchquad/tests/boole_test.py
@@ -67,6 +67,8 @@ def _run_boole_tests(backend, _precision):
         jit_integrate = None
 
         def integrate(*args, **kwargs):
+            # this function initializes the jit_integrate variable with a jit'ed integrate function
+            # which is then re-used on all other integrations (as is the point of JIT).
             nonlocal jit_integrate
             if jit_integrate is None:
                 jit_integrate = bl.get_jit_compiled_integrate(
@@ -95,14 +97,6 @@ def integrate(*args, **kwargs):
             None  # set to None again so can be re-used with new integrand shape
         )
 
-        def integrate(*args, **kwargs):
-            nonlocal jit_integrate
-            if jit_integrate is None:
-                jit_integrate = bl.get_jit_compiled_integrate(
-                    dim=1, N=N, backend=backend
-                )
-            return jit_integrate(*args, **kwargs)
-
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
diff --git a/torchquad/tests/monte_carlo_test.py b/torchquad/tests/monte_carlo_test.py
index 209cb550..3cb616de 100644
--- a/torchquad/tests/monte_carlo_test.py
+++ b/torchquad/tests/monte_carlo_test.py
@@ -88,6 +88,8 @@ def _run_monte_carlo_tests(backend, _precision):
         jit_integrate = None
 
         def integrate(*args, **kwargs):
+            # this function initializes the jit_integrate variable with a jit'ed integrate function
+            # which is then re-used on all other integrations (as is the point of JIT).
             nonlocal jit_integrate
             if jit_integrate is None:
                 jit_integrate = mc.get_jit_compiled_integrate(
@@ -127,14 +129,6 @@ def integrate(*args, **kwargs):
             None  # set to None again so can be re-used with new integrand shape
         )
 
-        def integrate(*args, **kwargs):
-            nonlocal jit_integrate
-            if jit_integrate is None:
-                jit_integrate = mc.get_jit_compiled_integrate(
-                    dim=1, N=N, backend=backend
-                )
-            return jit_integrate(*args, **kwargs)
-
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
diff --git a/torchquad/tests/simpson_test.py b/torchquad/tests/simpson_test.py
index c6dbe3c4..b9865ac8 100644
--- a/torchquad/tests/simpson_test.py
+++ b/torchquad/tests/simpson_test.py
@@ -89,11 +89,14 @@ def _run_simpson_tests(backend, _precision):
     for error in errors:
         assert error < 5e-9
 
+    # JIT Tests
     if backend != "numpy":
         N = 100001
         jit_integrate = None
 
         def integrate(*args, **kwargs):
+            # this function initializes the jit_integrate variable with a jit'ed integrate function
+            # which is then re-used on all other integrations (as is the point of JIT).
             nonlocal jit_integrate
             if jit_integrate is None:
                 jit_integrate = simp.get_jit_compiled_integrate(
@@ -120,15 +123,9 @@ def integrate(*args, **kwargs):
         for error in errors:
             assert error < 1e-7
 
-        jit_integrate = None
-
-        def integrate(*args, **kwargs):
-            nonlocal jit_integrate
-            if jit_integrate is None:
-                jit_integrate = simp.get_jit_compiled_integrate(
-                    dim=1, N=N, backend=backend
-                )
-            return jit_integrate(*args, **kwargs)
+        jit_integrate = (
+            None  # set to None again so can be re-used with new integrand shape
+        )
 
         errors, funcs = compute_integration_test_errors(
             integrate,
diff --git a/torchquad/tests/trapezoid_test.py b/torchquad/tests/trapezoid_test.py
index 11116253..086cbdb7 100644
--- a/torchquad/tests/trapezoid_test.py
+++ b/torchquad/tests/trapezoid_test.py
@@ -86,11 +86,14 @@ def _run_trapezoid_tests(backend, _precision):
     for error in errors:
         assert error < 7000
 
+    # JIT Tests
     if backend != "numpy":
         N = 100000
         jit_integrate = None
 
         def integrate(*args, **kwargs):
+            # this function initializes the jit_integrate variable with a jit'ed integrate function
+            # which is then re-used on all other integrations (as is the point of JIT).
             nonlocal jit_integrate
             if jit_integrate is None:
                 jit_integrate = tp.get_jit_compiled_integrate(
@@ -117,15 +120,9 @@ def integrate(*args, **kwargs):
         for error in errors:
             assert error < 1e-5
 
-        jit_integrate = None
-
-        def integrate(*args, **kwargs):
-            nonlocal jit_integrate
-            if jit_integrate is None:
-                jit_integrate = tp.get_jit_compiled_integrate(
-                    dim=1, N=N, backend=backend
-                )
-            return jit_integrate(*args, **kwargs)
+        jit_integrate = (
+            None  # set to None again so can be re-used with new integrand shape
+        )
 
         errors, funcs = compute_integration_test_errors(
             integrate,

From 586d256bdf4e10b0d85a879a05877df1e83a3482 Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Wed, 8 Mar 2023 11:24:21 +0100
Subject: [PATCH 12/21] (fix): increase test tolerance

---
 torchquad/tests/monte_carlo_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchquad/tests/monte_carlo_test.py b/torchquad/tests/monte_carlo_test.py
index 3cb616de..a8b48896 100644
--- a/torchquad/tests/monte_carlo_test.py
+++ b/torchquad/tests/monte_carlo_test.py
@@ -123,7 +123,7 @@ def integrate(*args, **kwargs):
             assert error < 2e-2
 
         for error in errors[10:]:
-            assert error < 28.03
+            assert error < 35.0
 
         jit_integrate = (
             None  # set to None again so can be re-used with new integrand shape

From d72936a6eb4e5324cd6b4f05a9182cd786d5fd44 Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Mon, 6 Mar 2023 16:06:12 +0100
Subject: [PATCH 13/21] (feat): add newton-cotest integrator tests

---
 torchquad/tests/boole_test.py     | 21 +++++++++++++++++++++
 torchquad/tests/simpson_test.py   | 20 ++++++++++++++++++++
 torchquad/tests/trapezoid_test.py | 20 ++++++++++++++++++++
 3 files changed, 61 insertions(+)

diff --git a/torchquad/tests/boole_test.py b/torchquad/tests/boole_test.py
index beeeace3..810ae78b 100644
--- a/torchquad/tests/boole_test.py
+++ b/torchquad/tests/boole_test.py
@@ -61,6 +61,27 @@ def _run_boole_tests(backend, _precision):
     # for error in errors:
     # assert error < 5e-9
 
+    # JIT Tests
+    if backend != "numpy":
+        N = 401
+
+        integrate = bl.get_jit_compiled_integrate(
+            dim=1, N=N, backend=backend
+        )
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+        )
+        print(f"1D Boole JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+        # Polynomials up to degree 5 can be integrated almost exactly with Boole.
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 5 or err < 6.33e-11
+        for error in errors:
+            assert error < 6.33e-11
+
 
 test_integrate_numpy = setup_test_for_backend(_run_boole_tests, "numpy", "float64")
 test_integrate_torch = setup_test_for_backend(_run_boole_tests, "torch", "float64")
diff --git a/torchquad/tests/simpson_test.py b/torchquad/tests/simpson_test.py
index 28172c56..23e3afe0 100644
--- a/torchquad/tests/simpson_test.py
+++ b/torchquad/tests/simpson_test.py
@@ -89,6 +89,26 @@ def _run_simpson_tests(backend, _precision):
     for error in errors:
         assert error < 5e-9
 
+    if backend != "numpy":
+        N = 401
+
+        integrate = simp.get_jit_compiled_integrate(
+            dim=1, N=N, backend=backend
+        )
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+        )
+        print(f"1D Simpson JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+        # Polynomials up to degree 5 can be integrated almost exactly with Boole.
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 5 or err < 6.33e-11
+        for error in errors:
+            assert error < 6.33e-11
+
 
 test_integrate_numpy = setup_test_for_backend(_run_simpson_tests, "numpy", "float64")
 test_integrate_torch = setup_test_for_backend(_run_simpson_tests, "torch", "float64")
diff --git a/torchquad/tests/trapezoid_test.py b/torchquad/tests/trapezoid_test.py
index 0d224d1a..c863a5fc 100644
--- a/torchquad/tests/trapezoid_test.py
+++ b/torchquad/tests/trapezoid_test.py
@@ -86,6 +86,26 @@ def _run_trapezoid_tests(backend, _precision):
     for error in errors:
         assert error < 7000
 
+    if backend != "numpy":
+        N = 401
+
+        integrate = tp.get_jit_compiled_integrate(
+            dim=1, N=N, backend=backend
+        )
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+        )
+        print(f"1D Trapezoid JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+        # Polynomials up to degree 5 can be integrated almost exactly with Boole.
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 5 or err < 6.33e-11
+        for error in errors:
+            assert error < 6.33e-11
+
 
 test_integrate_numpy = setup_test_for_backend(_run_trapezoid_tests, "numpy", "float64")
 test_integrate_torch = setup_test_for_backend(_run_trapezoid_tests, "torch", "float64")

From cb21a747aa94a3ff44febfceaba07b1f29f77f86 Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Mon, 6 Mar 2023 16:26:29 +0100
Subject: [PATCH 14/21] (fix): re-compile per integrand because of different
 shapes

---
 torchquad/tests/boole_test.py     |  9 +++++----
 torchquad/tests/simpson_test.py   | 16 ++++++++++------
 torchquad/tests/trapezoid_test.py | 16 ++++++++++------
 3 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/torchquad/tests/boole_test.py b/torchquad/tests/boole_test.py
index 810ae78b..35dae818 100644
--- a/torchquad/tests/boole_test.py
+++ b/torchquad/tests/boole_test.py
@@ -64,10 +64,11 @@ def _run_boole_tests(backend, _precision):
     # JIT Tests
     if backend != "numpy":
         N = 401
-
-        integrate = bl.get_jit_compiled_integrate(
-            dim=1, N=N, backend=backend
-        )
+        def integrate(*args, **kwargs):
+            jit_integrate = bl.get_jit_compiled_integrate(
+                dim=1, N=N, backend=backend
+            )
+            return jit_integrate(*args, **kwargs)
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
diff --git a/torchquad/tests/simpson_test.py b/torchquad/tests/simpson_test.py
index 23e3afe0..e5208ff7 100644
--- a/torchquad/tests/simpson_test.py
+++ b/torchquad/tests/simpson_test.py
@@ -90,11 +90,13 @@ def _run_simpson_tests(backend, _precision):
         assert error < 5e-9
 
     if backend != "numpy":
-        N = 401
+        N = 100001
 
-        integrate = simp.get_jit_compiled_integrate(
-            dim=1, N=N, backend=backend
-        )
+        def integrate(*args, **kwargs):
+            jit_integrate = simp.get_jit_compiled_integrate(
+                dim=1, N=N, backend=backend
+            )
+            return jit_integrate(*args, **kwargs)
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
@@ -105,9 +107,11 @@ def _run_simpson_tests(backend, _precision):
         print(f"1D Simpson JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
         # Polynomials up to degree 5 can be integrated almost exactly with Boole.
         for err, test_function in zip(errors, funcs):
-            assert test_function.get_order() > 5 or err < 6.33e-11
+            assert test_function.get_order() > 3 or (
+                err < 3e-11 if test_function.is_integrand_1d else err < 6e-10
+            )  # errors add up if the integrand is higher dimensional
         for error in errors:
-            assert error < 6.33e-11
+            assert error < 1e-7
 
 
 test_integrate_numpy = setup_test_for_backend(_run_simpson_tests, "numpy", "float64")
diff --git a/torchquad/tests/trapezoid_test.py b/torchquad/tests/trapezoid_test.py
index c863a5fc..f466f651 100644
--- a/torchquad/tests/trapezoid_test.py
+++ b/torchquad/tests/trapezoid_test.py
@@ -87,11 +87,13 @@ def _run_trapezoid_tests(backend, _precision):
         assert error < 7000
 
     if backend != "numpy":
-        N = 401
+        N = 100000
 
-        integrate = tp.get_jit_compiled_integrate(
-            dim=1, N=N, backend=backend
-        )
+        def integrate(*args, **kwargs):
+            jit_integrate = tp.get_jit_compiled_integrate(
+                dim=1, N=N, backend=backend
+            )
+            return jit_integrate(*args, **kwargs)
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
@@ -102,9 +104,11 @@ def _run_trapezoid_tests(backend, _precision):
         print(f"1D Trapezoid JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
         # Polynomials up to degree 5 can be integrated almost exactly with Boole.
         for err, test_function in zip(errors, funcs):
-            assert test_function.get_order() > 5 or err < 6.33e-11
+            assert test_function.get_order() > 1 or (
+                err < 2e-11 if test_function.is_integrand_1d else err < 5e-10
+            )  # errors add up if the integrand is higher dimensional
         for error in errors:
-            assert error < 6.33e-11
+            assert error < 1e-5
 
 
 test_integrate_numpy = setup_test_for_backend(_run_trapezoid_tests, "numpy", "float64")

From cc3daa5affbede15230a3e0a5be2c4970607b062 Mon Sep 17 00:00:00 2001
From: htoftevaag <htoftevaag@users.noreply.github.com>
Date: Mon, 6 Mar 2023 15:07:44 +0000
Subject: [PATCH 15/21] WORKFLOW: Format Python code with black

---
 torchquad/tests/boole_test.py     | 6 +++---
 torchquad/tests/simpson_test.py   | 9 +++++----
 torchquad/tests/trapezoid_test.py | 9 +++++----
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/torchquad/tests/boole_test.py b/torchquad/tests/boole_test.py
index 35dae818..50abb846 100644
--- a/torchquad/tests/boole_test.py
+++ b/torchquad/tests/boole_test.py
@@ -64,11 +64,11 @@ def _run_boole_tests(backend, _precision):
     # JIT Tests
     if backend != "numpy":
         N = 401
+
         def integrate(*args, **kwargs):
-            jit_integrate = bl.get_jit_compiled_integrate(
-                dim=1, N=N, backend=backend
-            )
+            jit_integrate = bl.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
             return jit_integrate(*args, **kwargs)
+
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
diff --git a/torchquad/tests/simpson_test.py b/torchquad/tests/simpson_test.py
index e5208ff7..398b4c42 100644
--- a/torchquad/tests/simpson_test.py
+++ b/torchquad/tests/simpson_test.py
@@ -93,10 +93,9 @@ def _run_simpson_tests(backend, _precision):
         N = 100001
 
         def integrate(*args, **kwargs):
-            jit_integrate = simp.get_jit_compiled_integrate(
-                dim=1, N=N, backend=backend
-            )
+            jit_integrate = simp.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
             return jit_integrate(*args, **kwargs)
+
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
@@ -104,7 +103,9 @@ def integrate(*args, **kwargs):
             use_complex=True,
             backend=backend,
         )
-        print(f"1D Simpson JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+        print(
+            f"1D Simpson JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}"
+        )
         # Polynomials up to degree 5 can be integrated almost exactly with Boole.
         for err, test_function in zip(errors, funcs):
             assert test_function.get_order() > 3 or (
diff --git a/torchquad/tests/trapezoid_test.py b/torchquad/tests/trapezoid_test.py
index f466f651..0d869dd1 100644
--- a/torchquad/tests/trapezoid_test.py
+++ b/torchquad/tests/trapezoid_test.py
@@ -90,10 +90,9 @@ def _run_trapezoid_tests(backend, _precision):
         N = 100000
 
         def integrate(*args, **kwargs):
-            jit_integrate = tp.get_jit_compiled_integrate(
-                dim=1, N=N, backend=backend
-            )
+            jit_integrate = tp.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
             return jit_integrate(*args, **kwargs)
+
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
@@ -101,7 +100,9 @@ def integrate(*args, **kwargs):
             use_complex=True,
             backend=backend,
         )
-        print(f"1D Trapezoid JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+        print(
+            f"1D Trapezoid JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}"
+        )
         # Polynomials up to degree 5 can be integrated almost exactly with Boole.
         for err, test_function in zip(errors, funcs):
             assert test_function.get_order() > 1 or (

From a51711909abbc04c28094b40ada21d60d5a3a002 Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Mon, 6 Mar 2023 16:51:50 +0100
Subject: [PATCH 16/21] (chore): remove erroneous boole comments

---
 torchquad/tests/simpson_test.py   | 1 -
 torchquad/tests/trapezoid_test.py | 1 -
 2 files changed, 2 deletions(-)

diff --git a/torchquad/tests/simpson_test.py b/torchquad/tests/simpson_test.py
index 398b4c42..629906fc 100644
--- a/torchquad/tests/simpson_test.py
+++ b/torchquad/tests/simpson_test.py
@@ -106,7 +106,6 @@ def integrate(*args, **kwargs):
         print(
             f"1D Simpson JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}"
         )
-        # Polynomials up to degree 5 can be integrated almost exactly with Boole.
         for err, test_function in zip(errors, funcs):
             assert test_function.get_order() > 3 or (
                 err < 3e-11 if test_function.is_integrand_1d else err < 6e-10
diff --git a/torchquad/tests/trapezoid_test.py b/torchquad/tests/trapezoid_test.py
index 0d869dd1..e9853ac6 100644
--- a/torchquad/tests/trapezoid_test.py
+++ b/torchquad/tests/trapezoid_test.py
@@ -103,7 +103,6 @@ def integrate(*args, **kwargs):
         print(
             f"1D Trapezoid JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}"
         )
-        # Polynomials up to degree 5 can be integrated almost exactly with Boole.
         for err, test_function in zip(errors, funcs):
             assert test_function.get_order() > 1 or (
                 err < 2e-11 if test_function.is_integrand_1d else err < 5e-10

From 76278dfa92d476bcefdbe181e59d96e48db9289d Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Mon, 6 Mar 2023 16:55:10 +0100
Subject: [PATCH 17/21] (fix): add mc test

---
 torchquad/tests/monte_carlo_test.py | 33 +++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/torchquad/tests/monte_carlo_test.py b/torchquad/tests/monte_carlo_test.py
index ab624bb0..8788d292 100644
--- a/torchquad/tests/monte_carlo_test.py
+++ b/torchquad/tests/monte_carlo_test.py
@@ -82,6 +82,39 @@ def _run_monte_carlo_tests(backend, _precision):
     for error in errors:
         assert error < 26
 
+    # JIT Tests
+    if backend != "numpy":
+        N = 100000
+
+        def integrate(*args, **kwargs):
+            jit_integrate = mc.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
+            return jit_integrate(*args, **kwargs)
+
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+        )
+        print(f"1D MC JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
+
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 0 or err < 1e-14
+
+        # If this breaks check if test functions in helper_functions changed.
+        for error in errors[:3]:
+            assert error < 7e-3
+
+        assert errors[3] < 0.5
+        assert errors[4] < 32.0
+
+        for error in errors[6:10]:
+            assert error < 2e-2
+
+        for error in errors[10:]:
+            assert error < 28.03
+
 
 test_integrate_numpy = setup_test_for_backend(
     _run_monte_carlo_tests, "numpy", "float32"

From 669a0a25a72b7b646008e18bc5b5669093aecff6 Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Mon, 6 Mar 2023 17:02:48 +0100
Subject: [PATCH 18/21] (chore): add note about `integrand` shape

---
 docs/source/tutorial.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/source/tutorial.rst b/docs/source/tutorial.rst
index d2164e21..c9d034d1 100644
--- a/docs/source/tutorial.rst
+++ b/docs/source/tutorial.rst
@@ -586,7 +586,7 @@ Compiling the integrate method
 ``````````````````````````````
 
 To speed up the quadrature in situations where it is executed often with the
-same number of points ``N`` and dimensionality ``dim``,
+same number of points ``N``, dimensionality ``dim``, and shape of the ``integrand``,
 we can JIT-compile the performance-relevant parts of the integrate method:
 
 .. code:: ipython3

From dd1bc5d3d034e1d406011572854657bf6bda8d2b Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Tue, 7 Mar 2023 14:07:00 +0100
Subject: [PATCH 19/21] (feat): switch to re-use fo the jit integral

---
 torchquad/tests/boole_test.py       | 41 +++++++++++++++++++++++++++--
 torchquad/tests/helper_functions.py |  8 +++---
 torchquad/tests/monte_carlo_test.py | 41 ++++++++++++++++++++++++++---
 torchquad/tests/simpson_test.py     | 37 +++++++++++++++++++++++++-
 torchquad/tests/trapezoid_test.py   | 39 +++++++++++++++++++++++++--
 5 files changed, 155 insertions(+), 11 deletions(-)

diff --git a/torchquad/tests/boole_test.py b/torchquad/tests/boole_test.py
index 50abb846..a2025543 100644
--- a/torchquad/tests/boole_test.py
+++ b/torchquad/tests/boole_test.py
@@ -64,9 +64,14 @@ def _run_boole_tests(backend, _precision):
     # JIT Tests
     if backend != "numpy":
         N = 401
+        jit_integrate = None
 
         def integrate(*args, **kwargs):
-            jit_integrate = bl.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = bl.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
             return jit_integrate(*args, **kwargs)
 
         errors, funcs = compute_integration_test_errors(
@@ -75,8 +80,40 @@ def integrate(*args, **kwargs):
             integration_dim=1,
             use_complex=True,
             backend=backend,
+            filter_test_functions=lambda x: x.is_integrand_1d,
+        )
+        print(
+            f"1D Boole JIT Test passed for 1D integrands. N: {N}, backend: {backend}, Errors: {errors}"
+        )
+        # Polynomials up to degree 5 can be integrated almost exactly with Boole.
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 5 or err < 6.33e-11
+        for error in errors:
+            assert error < 6.33e-11
+
+        jit_integrate = (
+            None  # set to None again so can be re-used with new integrand shape
+        )
+
+        def integrate(*args, **kwargs):
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = bl.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
+            return jit_integrate(*args, **kwargs)
+
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+            filter_test_functions=lambda x: x.integrand_dims == [2, 2, 2],
+        )
+        print(
+            f"1D Boole JIT Test passed for [2, 2, 2] dimensional integrands. N: {N}, backend: {backend}, Errors: {errors}"
         )
-        print(f"1D Boole JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
         # Polynomials up to degree 5 can be integrated almost exactly with Boole.
         for err, test_function in zip(errors, funcs):
             assert test_function.get_order() > 5 or err < 6.33e-11
diff --git a/torchquad/tests/helper_functions.py b/torchquad/tests/helper_functions.py
index c0dcc884..88a8ee84 100644
--- a/torchquad/tests/helper_functions.py
+++ b/torchquad/tests/helper_functions.py
@@ -324,6 +324,7 @@ def compute_integration_test_errors(
     use_complex,
     backend,
     use_multi_dim_integrand=True,
+    filter_test_functions=lambda x: x,
 ):
     """Computes errors on all test functions for given dimension and integrator.
 
@@ -334,7 +335,7 @@ def compute_integration_test_errors(
         use_complex (Boolean): If True, skip complex example functions.
         backend (string): Numerical backend for the example functions.
         use_multi_dim_integrand (bool, optional): Whether or not to allow for a multi-dimensional integrand i.e an array of integrands
-
+        filter_test_functions (function, optional): function for filtering which test functions to run
     Returns:
         (list, list): Absolute errors on all example functions and the chosen
             example functions
@@ -344,8 +345,9 @@ def compute_integration_test_errors(
 
     # Compute integration errors on the chosen functions and remember those
     # functions
-    for test_function in get_test_functions(
-        integration_dim, backend, use_multi_dim_integrand
+    for test_function in filter(
+        filter_test_functions,
+        get_test_functions(integration_dim, backend, use_multi_dim_integrand),
     ):
         if not use_complex and test_function.is_complex:
             continue
diff --git a/torchquad/tests/monte_carlo_test.py b/torchquad/tests/monte_carlo_test.py
index 8788d292..209cb550 100644
--- a/torchquad/tests/monte_carlo_test.py
+++ b/torchquad/tests/monte_carlo_test.py
@@ -85,9 +85,14 @@ def _run_monte_carlo_tests(backend, _precision):
     # JIT Tests
     if backend != "numpy":
         N = 100000
+        jit_integrate = None
 
         def integrate(*args, **kwargs):
-            jit_integrate = mc.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = mc.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
             return jit_integrate(*args, **kwargs)
 
         errors, funcs = compute_integration_test_errors(
@@ -96,15 +101,18 @@ def integrate(*args, **kwargs):
             integration_dim=1,
             use_complex=True,
             backend=backend,
+            filter_test_functions=lambda x: x.is_integrand_1d,
+        )
+        print(
+            f"1D MC JIT Test passed for 1D integrands. N: {N}, backend: {backend}, Errors: {errors}"
         )
-        print(f"1D MC JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}")
 
         for err, test_function in zip(errors, funcs):
             assert test_function.get_order() > 0 or err < 1e-14
 
         # If this breaks check if test functions in helper_functions changed.
         for error in errors[:3]:
-            assert error < 7e-3
+            assert error < 1e-2
 
         assert errors[3] < 0.5
         assert errors[4] < 32.0
@@ -115,6 +123,33 @@ def integrate(*args, **kwargs):
         for error in errors[10:]:
             assert error < 28.03
 
+        jit_integrate = (
+            None  # set to None again so can be re-used with new integrand shape
+        )
+
+        def integrate(*args, **kwargs):
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = mc.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
+            return jit_integrate(*args, **kwargs)
+
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+            filter_test_functions=lambda x: x.integrand_dims == [2, 2, 2],
+        )
+        print(
+            f"1D MC JIT Test passed for [2, 2, 2] dimensional integrands. N: {N}, backend: {backend}, Errors: {errors}"
+        )
+
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 0 or err < 1e-14
+
 
 test_integrate_numpy = setup_test_for_backend(
     _run_monte_carlo_tests, "numpy", "float32"
diff --git a/torchquad/tests/simpson_test.py b/torchquad/tests/simpson_test.py
index 629906fc..c6dbe3c4 100644
--- a/torchquad/tests/simpson_test.py
+++ b/torchquad/tests/simpson_test.py
@@ -91,9 +91,14 @@ def _run_simpson_tests(backend, _precision):
 
     if backend != "numpy":
         N = 100001
+        jit_integrate = None
 
         def integrate(*args, **kwargs):
-            jit_integrate = simp.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = simp.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
             return jit_integrate(*args, **kwargs)
 
         errors, funcs = compute_integration_test_errors(
@@ -102,7 +107,9 @@ def integrate(*args, **kwargs):
             integration_dim=1,
             use_complex=True,
             backend=backend,
+            filter_test_functions=lambda x: x.is_integrand_1d,
         )
+
         print(
             f"1D Simpson JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}"
         )
@@ -113,6 +120,34 @@ def integrate(*args, **kwargs):
         for error in errors:
             assert error < 1e-7
 
+        jit_integrate = None
+
+        def integrate(*args, **kwargs):
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = simp.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
+            return jit_integrate(*args, **kwargs)
+
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+            filter_test_functions=lambda x: x.integrand_dims == [2, 2, 2],
+        )
+        print(
+            f"1D Simpson JIT Test passed for [2, 2, 2] dimensional integrands. N: {N}, backend: {backend}, Errors: {errors}"
+        )
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 3 or (
+                err < 3e-11 if test_function.is_integrand_1d else err < 6e-10
+            )  # errors add up if the integrand is higher dimensional
+        for error in errors:
+            assert error < 1e-7
+
 
 test_integrate_numpy = setup_test_for_backend(_run_simpson_tests, "numpy", "float64")
 test_integrate_torch = setup_test_for_backend(_run_simpson_tests, "torch", "float64")
diff --git a/torchquad/tests/trapezoid_test.py b/torchquad/tests/trapezoid_test.py
index e9853ac6..11116253 100644
--- a/torchquad/tests/trapezoid_test.py
+++ b/torchquad/tests/trapezoid_test.py
@@ -88,9 +88,14 @@ def _run_trapezoid_tests(backend, _precision):
 
     if backend != "numpy":
         N = 100000
+        jit_integrate = None
 
         def integrate(*args, **kwargs):
-            jit_integrate = tp.get_jit_compiled_integrate(dim=1, N=N, backend=backend)
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = tp.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
             return jit_integrate(*args, **kwargs)
 
         errors, funcs = compute_integration_test_errors(
@@ -99,9 +104,39 @@ def integrate(*args, **kwargs):
             integration_dim=1,
             use_complex=True,
             backend=backend,
+            filter_test_functions=lambda x: x.is_integrand_1d,
+        )
+
+        print(
+            f"1D Trapezoid JIT Test passed for 1D integrands. N: {N}, backend: {backend}, Errors: {errors}"
+        )
+        for err, test_function in zip(errors, funcs):
+            assert test_function.get_order() > 1 or (
+                err < 2e-11 if test_function.is_integrand_1d else err < 5e-10
+            )  # errors add up if the integrand is higher dimensional
+        for error in errors:
+            assert error < 1e-5
+
+        jit_integrate = None
+
+        def integrate(*args, **kwargs):
+            nonlocal jit_integrate
+            if jit_integrate is None:
+                jit_integrate = tp.get_jit_compiled_integrate(
+                    dim=1, N=N, backend=backend
+                )
+            return jit_integrate(*args, **kwargs)
+
+        errors, funcs = compute_integration_test_errors(
+            integrate,
+            {},
+            integration_dim=1,
+            use_complex=True,
+            backend=backend,
+            filter_test_functions=lambda x: x.integrand_dims == [2, 2, 2],
         )
         print(
-            f"1D Trapezoid JIT Test passed. N: {N}, backend: {backend}, Errors: {errors}"
+            f"1D Trapezoid JIT Test passed for [2, 2, 2] dimensional integrands. N: {N}, backend: {backend}, Errors: {errors}"
         )
         for err, test_function in zip(errors, funcs):
             assert test_function.get_order() > 1 or (

From 37da0ed64e2584c66702b5ba684aa1ae95bf2857 Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Wed, 8 Mar 2023 10:54:06 +0100
Subject: [PATCH 20/21] (chore): address comments

---
 torchquad/tests/boole_test.py       | 10 ++--------
 torchquad/tests/monte_carlo_test.py | 10 ++--------
 torchquad/tests/simpson_test.py     | 15 ++++++---------
 torchquad/tests/trapezoid_test.py   | 15 ++++++---------
 4 files changed, 16 insertions(+), 34 deletions(-)

diff --git a/torchquad/tests/boole_test.py b/torchquad/tests/boole_test.py
index a2025543..aaea3f1a 100644
--- a/torchquad/tests/boole_test.py
+++ b/torchquad/tests/boole_test.py
@@ -67,6 +67,8 @@ def _run_boole_tests(backend, _precision):
         jit_integrate = None
 
         def integrate(*args, **kwargs):
+            # this function initializes the jit_integrate variable with a jit'ed integrate function
+            # which is then re-used on all other integrations (as is the point of JIT).
             nonlocal jit_integrate
             if jit_integrate is None:
                 jit_integrate = bl.get_jit_compiled_integrate(
@@ -95,14 +97,6 @@ def integrate(*args, **kwargs):
             None  # set to None again so can be re-used with new integrand shape
         )
 
-        def integrate(*args, **kwargs):
-            nonlocal jit_integrate
-            if jit_integrate is None:
-                jit_integrate = bl.get_jit_compiled_integrate(
-                    dim=1, N=N, backend=backend
-                )
-            return jit_integrate(*args, **kwargs)
-
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
diff --git a/torchquad/tests/monte_carlo_test.py b/torchquad/tests/monte_carlo_test.py
index 209cb550..3cb616de 100644
--- a/torchquad/tests/monte_carlo_test.py
+++ b/torchquad/tests/monte_carlo_test.py
@@ -88,6 +88,8 @@ def _run_monte_carlo_tests(backend, _precision):
         jit_integrate = None
 
         def integrate(*args, **kwargs):
+            # this function initializes the jit_integrate variable with a jit'ed integrate function
+            # which is then re-used on all other integrations (as is the point of JIT).
             nonlocal jit_integrate
             if jit_integrate is None:
                 jit_integrate = mc.get_jit_compiled_integrate(
@@ -127,14 +129,6 @@ def integrate(*args, **kwargs):
             None  # set to None again so can be re-used with new integrand shape
         )
 
-        def integrate(*args, **kwargs):
-            nonlocal jit_integrate
-            if jit_integrate is None:
-                jit_integrate = mc.get_jit_compiled_integrate(
-                    dim=1, N=N, backend=backend
-                )
-            return jit_integrate(*args, **kwargs)
-
         errors, funcs = compute_integration_test_errors(
             integrate,
             {},
diff --git a/torchquad/tests/simpson_test.py b/torchquad/tests/simpson_test.py
index c6dbe3c4..b9865ac8 100644
--- a/torchquad/tests/simpson_test.py
+++ b/torchquad/tests/simpson_test.py
@@ -89,11 +89,14 @@ def _run_simpson_tests(backend, _precision):
     for error in errors:
         assert error < 5e-9
 
+    # JIT Tests
     if backend != "numpy":
         N = 100001
         jit_integrate = None
 
         def integrate(*args, **kwargs):
+            # this function initializes the jit_integrate variable with a jit'ed integrate function
+            # which is then re-used on all other integrations (as is the point of JIT).
             nonlocal jit_integrate
             if jit_integrate is None:
                 jit_integrate = simp.get_jit_compiled_integrate(
@@ -120,15 +123,9 @@ def integrate(*args, **kwargs):
         for error in errors:
             assert error < 1e-7
 
-        jit_integrate = None
-
-        def integrate(*args, **kwargs):
-            nonlocal jit_integrate
-            if jit_integrate is None:
-                jit_integrate = simp.get_jit_compiled_integrate(
-                    dim=1, N=N, backend=backend
-                )
-            return jit_integrate(*args, **kwargs)
+        jit_integrate = (
+            None  # set to None again so can be re-used with new integrand shape
+        )
 
         errors, funcs = compute_integration_test_errors(
             integrate,
diff --git a/torchquad/tests/trapezoid_test.py b/torchquad/tests/trapezoid_test.py
index 11116253..086cbdb7 100644
--- a/torchquad/tests/trapezoid_test.py
+++ b/torchquad/tests/trapezoid_test.py
@@ -86,11 +86,14 @@ def _run_trapezoid_tests(backend, _precision):
     for error in errors:
         assert error < 7000
 
+    # JIT Tests
     if backend != "numpy":
         N = 100000
         jit_integrate = None
 
         def integrate(*args, **kwargs):
+            # this function initializes the jit_integrate variable with a jit'ed integrate function
+            # which is then re-used on all other integrations (as is the point of JIT).
             nonlocal jit_integrate
             if jit_integrate is None:
                 jit_integrate = tp.get_jit_compiled_integrate(
@@ -117,15 +120,9 @@ def integrate(*args, **kwargs):
         for error in errors:
             assert error < 1e-5
 
-        jit_integrate = None
-
-        def integrate(*args, **kwargs):
-            nonlocal jit_integrate
-            if jit_integrate is None:
-                jit_integrate = tp.get_jit_compiled_integrate(
-                    dim=1, N=N, backend=backend
-                )
-            return jit_integrate(*args, **kwargs)
+        jit_integrate = (
+            None  # set to None again so can be re-used with new integrand shape
+        )
 
         errors, funcs = compute_integration_test_errors(
             integrate,

From 155e235ae96a6c73385418e37268d51bcd5514fb Mon Sep 17 00:00:00 2001
From: ilan-gold <ilanbassgold@gmail.com>
Date: Wed, 8 Mar 2023 11:24:21 +0100
Subject: [PATCH 21/21] (fix): increase test tolerance

---
 torchquad/tests/monte_carlo_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchquad/tests/monte_carlo_test.py b/torchquad/tests/monte_carlo_test.py
index 3cb616de..a8b48896 100644
--- a/torchquad/tests/monte_carlo_test.py
+++ b/torchquad/tests/monte_carlo_test.py
@@ -123,7 +123,7 @@ def integrate(*args, **kwargs):
             assert error < 2e-2
 
         for error in errors[10:]:
-            assert error < 28.03
+            assert error < 35.0
 
         jit_integrate = (
             None  # set to None again so can be re-used with new integrand shape