(fix): use static_argnums for jax jit

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_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

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

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")