Memory leak issues with lax.scan

[ericmjl]

Putting up this issue so that I can keep track of progress.

This is related to google/jax#3348, in which the OP followed a very similar pattern to what we have used in unirep:

        fx = partial(apply_fun_scan, p1)
        _, ht_new = lax.scan(fx, p2, inputs)  

I suspect we have to get rid of this pattern in order to avoid memory leaks with JAX. (I think this is what might have crashed a colleague's machine at work.) I thought of an idea in google/jax#3348 (comment), so maybe that could be a starting point for me to kickstart.