Repository: f0uriest/quadax

Issue:
Quadrature within jitted functions can cause leaked trace
Hi! This library is fantastic for performing numerical quadrature with JAX!

I encountered an issue when enabling JAX’s leak-checking mechanism for tracers. The following example raises an exception:

```python
import jax
import jax.numpy as jnp
from quadax import quadgk

@jax.jit
def integral(interval):
    return quadgk(jnp.square, interval)

with jax.checking_leaks():
    jax.block_until_ready(integral([0.0, 1.0]))
```

The error message indicates a trace leak:

...... Exception: Leaked trace MainTrace(1,DynamicJaxprTrace). Leaked tracer(s):
Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>
This DynamicJaxprTracer was created on line .../python3.9/site-packages/quadax/utils.py:99 (map_interval)
<DynamicJaxprTracer 5063988464> is referred to by ......

Hints from issue comments:
@f0uriest: Hmm. I'm able to reproduce this but haven't been able to fully isolate it, so still working on a fix.@f0uriest: OK, think I figured it out. Should be resolved by #48 

Solution patch:
diff --git a/.github/dependabot.yml b/.github/dependabot.yml
index 60841fd..f4c9d40 100644
--- a/.github/dependabot.yml
+++ b/.github/dependabot.yml
@@ -6,7 +6,7 @@ updates:
   - package-ecosystem: "pip"
     directory: "/" # Location of package manifests
     schedule:
-      interval: "weekly"
+      interval: "daily"
   - package-ecosystem: "github-actions"
     directory: "/"
     schedule:
diff --git a/.github/workflows/scheduled.yml b/.github/workflows/scheduled.yml
index f1b1c27..9ff2038 100644
--- a/.github/workflows/scheduled.yml
+++ b/.github/workflows/scheduled.yml
@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.9', '3.10', '3.11', '3.12']
+        python-version: ['3.10', '3.11', '3.12']
 
     steps:
       - uses: actions/checkout@v4
@@ -22,7 +22,8 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -r requirements-dev.txt
+          pip install --upgrade -r requirements-dev.txt
+          pip list
       - name: Test with pytest
         run: |
           pwd
diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml
index 1cb1612..dfa105d 100644
--- a/.github/workflows/unittest.yml
+++ b/.github/workflows/unittest.yml
@@ -23,7 +23,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.9', '3.10', '3.11', '3.12']
+        python-version: ['3.10', '3.11', '3.12']
 
     steps:
       - uses: actions/checkout@v4
@@ -34,7 +34,8 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -r requirements-dev.txt
+          pip install --upgrade -r requirements-dev.txt
+          pip list
       - name: Test with pytest
         run: |
           pwd
diff --git a/quadax/adaptive.py b/quadax/adaptive.py
index a374a9f..a32f518 100644
--- a/quadax/adaptive.py
+++ b/quadax/adaptive.py
@@ -2,6 +2,7 @@
 
 import warnings
 
+import equinox as eqx
 import jax
 import jax.numpy as jnp
 
@@ -28,6 +29,7 @@
 DIVERGENT = 5
 
 
+@eqx.filter_jit
 def quadgk(
     fun,
     interval,
@@ -128,6 +130,7 @@ def quadgk(
     return y, info
 
 
+@eqx.filter_jit
 def quadcc(
     fun,
     interval,
@@ -227,6 +230,7 @@ def quadcc(
     return y, info
 
 
+@eqx.filter_jit
 def quadts(
     fun,
     interval,
@@ -325,6 +329,7 @@ def quadts(
     return y, info
 
 
+@eqx.filter_jit
 def adaptive_quadrature(
     rule,
     fun,
diff --git a/quadax/romberg.py b/quadax/romberg.py
index 32c4092..876e90e 100644
--- a/quadax/romberg.py
+++ b/quadax/romberg.py
@@ -1,5 +1,6 @@
 """Romberg integration aka adaptive trapezoid with Richardson extrapolation."""
 
+import equinox as eqx
 import jax
 import jax.numpy as jnp
 
@@ -14,6 +15,7 @@
 )
 
 
+@eqx.filter_jit
 def romberg(
     fun,
     interval,
@@ -147,6 +149,7 @@ def mloop(m, result):
     return y, out
 
 
+@eqx.filter_jit
 def rombergts(
     fun,
     interval,
diff --git a/quadax/utils.py b/quadax/utils.py
index 42d0d14..254a39c 100644
--- a/quadax/utils.py
+++ b/quadax/utils.py
@@ -1,8 +1,8 @@
 """Utility functions for parsing inputs, mapping coordinates etc."""
 
-from functools import partial
-from typing import NamedTuple, Union
+from typing import Any, NamedTuple, Union
 
+import equinox as eqx
 import jax
 import jax.numpy as jnp
 
@@ -22,7 +22,7 @@ def _map_linear(t, a, b):
     c = (b - a) / 2
     d = (b + a) / 2
     x = d + c * t
-    w = c
+    w = c * jnp.ones_like(t)
     return x.squeeze(), w.squeeze()
 
 
@@ -73,6 +73,10 @@ def _map_ninfb_inv(x, a, b):
     return t.squeeze()
 
 
+MAPFUNS = [_map_linear, _map_ninfb, _map_ainf, _map_ninfinf]
+MAPFUNS_INV = [_map_linear_inv, _map_ninfb_inv, _map_ainf_inv, _map_ninfinf_inv]
+
+
 def map_interval(fun, interval):
     """Map a function over an arbitrary interval [a, b] to the interval [-1, 1].
 
@@ -110,22 +114,31 @@ def map_interval(fun, interval):
     # 2 : a finite, b = inf
     # 3 : both infinite
     bitmask = jnp.isinf(a) + 2 * jnp.isinf(b)
-    mapfuns = [_map_linear, _map_ninfb, _map_ainf, _map_ninfinf]
-    mapfuns_inv = [_map_linear_inv, _map_ninfb_inv, _map_ainf_inv, _map_ninfinf_inv]
-
-    @jax.jit
-    def fun_mapped(t, *args):
-        x, w = jax.lax.switch(bitmask, mapfuns, t, a, b)
-        return sgn * w * fun(x, *args)
 
+    fun_mapped = _MappedFunction(fun, bitmask, sgn, a, b)
     # map original breakpoints to new domain
-    interval_t = jax.lax.switch(bitmask, mapfuns_inv, interval, a, b)
+    interval_t = jax.lax.switch(bitmask, MAPFUNS_INV, interval, a, b)
     # +/-inf gets mapped to +/-1 but numerically evaluates to nan so we replace that.
     interval_t = jnp.where(interval == jnp.inf, 1, interval_t)
     interval_t = jnp.where(interval == -jnp.inf, -1, interval_t)
     return fun_mapped, interval_t
 
 
+class _MappedFunction(eqx.Module):
+    """Function mapped to unit interval [-1,1]."""
+
+    fun: callable
+    bitmask: int
+    sgn: int
+    a: float
+    b: float
+
+    @eqx.filter_jit
+    def __call__(self, t, *args):
+        x, w = jax.lax.switch(self.bitmask, MAPFUNS, t, self.a, self.b)
+        return self.sgn * w * self.fun(x, *args)
+
+
 def tanhsinh_transform(fun, interval):
     """Transform a function by mapping with tanh-sinh.
 
@@ -154,10 +167,7 @@ def tanhsinh_transform(fun, interval):
     # map a, b -> [-1, 1]
     fun, interval = map_interval(fun, interval)
 
-    # map [-1, 1] to [-inf, inf], but with mass concentrated near 0
-    xk = lambda t: jnp.tanh(jnp.pi / 2 * jnp.sinh(t))
-    wk = lambda t: jnp.pi / 2 * jnp.cosh(t) / jnp.cosh(jnp.pi / 2 * jnp.sinh(t)) ** 2
-    func = lambda t, *args: fun(xk(t), *args) * wk(t)
+    func = _TanhSinhTransformedFunction(fun)
 
     # we generally only need to integrate ~[-3, 3] or ~[-4, 4]
     # we don't want to include the endpoint that maps to x==1 to avoid
@@ -172,7 +182,26 @@ def get_tmax(xmax):
     # inverse of tanh-sinh transformation for x = 1-eps
     tmax = get_tmax(jnp.array(1.0) - 10 * jnp.finfo(jnp.array(1.0)).eps)
     interval_t = jnp.array([-tmax, tmax])
-    return jax.jit(func), interval_t
+    return func, interval_t
+
+
+# map [-1, 1] to [-inf, inf], but with mass concentrated near 0
+tanhsinh_x = lambda t: jnp.tanh(jnp.pi / 2 * jnp.sinh(t))
+tanhsinh_w = (
+    lambda t: jnp.pi / 2 * jnp.cosh(t) / jnp.cosh(jnp.pi / 2 * jnp.sinh(t)) ** 2
+)
+
+
+class _TanhSinhTransformedFunction(eqx.Module):
+    """Function transformed by tanh-sinh transformation."""
+
+    fun: callable
+
+    @eqx.filter_jit
+    def __call__(self, t, *args):
+        x = tanhsinh_x(t)
+        w = tanhsinh_w(t)
+        return self.fun(x, *args) * w
 
 
 messages = {
@@ -226,13 +255,24 @@ def wrap_func(fun, args):
     # need to make sure we get the correct shape for array valued integrands
     outsig = "(" + ",".join("n" + str(i) for i in range(len(f.shape))) + ")"
 
-    @jax.jit
-    @partial(jnp.vectorize, signature="()->" + outsig)
-    def wrapped(x):
-        f = fun(x, *args)
-        return jnp.where(jnp.isfinite(f), f, 0.0)
+    return _WrappedFunction(fun, args, outsig)
 
-    return wrapped
+
+class _WrappedFunction(eqx.Module):
+    """Wraps a function in jit/vectorize and masks out inf/nans."""
+
+    fun: callable
+    args: Any
+    outsig: str
+
+    @eqx.filter_jit
+    def __call__(self, x):
+        f = jnp.vectorize(
+            self.fun,
+            excluded=tuple(range(1, len(self.args) + 1)),
+            signature="()->" + self.outsig,
+        )(x, *self.args)
+        return jnp.where(jnp.isfinite(f), f, 0.0)
 
 
 class QuadratureInfo(NamedTuple):


Test patch:
diff --git a/tests/test_adaptive.py b/tests/test_adaptive.py
index f0b74c9..2d405b4 100644
--- a/tests/test_adaptive.py
+++ b/tests/test_adaptive.py
@@ -1,5 +1,6 @@
 """Tests for adaptive quadrature routines."""
 
+import jax
 import jax.numpy as jnp
 import numpy as np
 import pytest
@@ -706,3 +707,42 @@ def test_prob16(self):
         self._base(16, 1e-4)
         self._base(16, 1e-8)
         self._base(16, 1e-12)
+
+
+def test_escaped_tracers():
+    """Test that no tracers escape, related to gh issue 18."""
+
+    @jax.jit
+    def integral_quadgk(interval):
+        return quadgk(jnp.square, interval)
+
+    with jax.checking_leaks():
+        jax.block_until_ready(integral_quadgk([0.0, 1.0]))
+
+    @jax.jit
+    def integral_quadcc(interval):
+        return quadcc(jnp.square, interval)
+
+    with jax.checking_leaks():
+        jax.block_until_ready(integral_quadcc([0.0, 1.0]))
+
+    @jax.jit
+    def integral_quadts(interval):
+        return quadts(jnp.square, interval)
+
+    with jax.checking_leaks():
+        jax.block_until_ready(integral_quadts([0.0, 1.0]))
+
+    @jax.jit
+    def integral_romberg(interval):
+        return romberg(jnp.square, interval)
+
+    with jax.checking_leaks():
+        jax.block_until_ready(integral_romberg([0.0, 1.0]))
+
+    @jax.jit
+    def integral_rombergts(interval):
+        return rombergts(jnp.square, interval)
+
+    with jax.checking_leaks():
+        jax.block_until_ready(integral_rombergts([0.0, 1.0]))