Quadrature within jitted functions can cause leaked trace

[wcxve]

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:

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

[f0uriest]

Hmm. I'm able to reproduce this but haven't been able to fully isolate it, so still working on a fix.