Problem with jax_cosmo caching system

[ASKabalan]

@EiffL
I created this 'short' snippet to reproduce the leak error that I systematically get when I am doing forward modeling with jaxpm

import jax
import jax.numpy as jnp
from diffrax import Euler, ODETerm, diffeqsolve

import jax_cosmo as jc

cosmo = jc.Planck15()


def ode_with_chi(t, y, cosmo):
    distance = jc.background.radial_comoving_distance(cosmo, t).squeeze()
    return jnp.squeeze(distance) * 0.1 + y


def ode_with_gp(t, y, cosmo):
    growth_factor = jc.background.growth_factor(cosmo, jnp.array([t]))
    return jnp.squeeze(growth_factor) * 0.1 + y


def ode_wrapper(cosmo):
    def fn(t, y, args):
        return ode_with_chi(t, y, cosmo)

    return fn


def integrate(cosmo, y, ode_term):
    result = diffeqsolve(
        ode_term,
        Euler(),
        t0=0.0,
        t1=1.0,
        dt0=0.1,
        y0=y,
        args=cosmo,  # This is the key part that causes the leak!
    )
    return result.ys


@jax.jit
def integrate_with_chi(cosmo):
    ode_term = ODETerm(ode_with_chi)
    return integrate(cosmo, jnp.array([1.0]), ode_term)


@jax.jit
def integrate_with_gp(cosmo):
    ode_term = ODETerm(ode_with_gp)
    return integrate(cosmo, jnp.array([1.0]), ode_term)


@jax.jit
def integrate_with_wrapper(cosmo):
    ode_term = ODETerm(ode_wrapper(cosmo))
    return integrate(cosmo, jnp.array([1.0]), ode_term)


@jax.jit
def integrate_with_chi_precompute(cosmo):
    distance = jc.background.radial_comoving_distance(cosmo, 0.1).squeeze()
    y0 = jnp.array([distance])
    ode_term = ODETerm(ode_with_chi)
    return integrate(cosmo, y0, ode_term)


@jax.jit
def integrate_with_chi_precompute_and_wrapper(cosmo):
    _ = jc.background.radial_comoving_distance(cosmo, 0.1).squeeze()
    y0 = jnp.array([1.0])
    ode_term = ODETerm(ode_wrapper(cosmo))
    return integrate(cosmo, y0, ode_term)


integrate_with_chi(cosmo)  # Leaks
integrate_with_gp(cosmo)  # Leaks
integrate_with_wrapper(cosmo)  # Leaks
integrate_with_chi_precompute(cosmo)  # Leaks
integrate_with_chi_precompute_and_wrapper(cosmo)  # DOES NOT leak

I think I can even reproduce without diffrax I think since the error here is caused by the fact that diffrax calls eval_shape which traces the caching twice

leak report

  File "/home/wassim/micromamba/envs/jax/lib/python3.10/site-packages/diffrax/_integrate.py", line 621, in loop
    _, traced_jump, traced_result = eqx.filter_eval_shape(body_fun_aux, init_state)
  File "/home/wassim/micromamba/envs/jax/lib/python3.10/contextlib.py", line 142, in __exit__
    next(self.gen)
Exception: Leaked trace DynamicJaxprTrace. Leaked tracer(s):

Traced<ShapedArray(float32[256])>with<DynamicJaxprTrace>
The error occurred while tracing the function _fn at /home/wassim/micromamba/envs/jax/lib/python3.10/site-packages/equinox/_eval_shape.py:31 for jit. 
<DynamicJaxprTracer 130500598565440> is referred to by <dict 130500600235072>['a']
<dict 130500600235072> is referred to by <dict 130500599962048>['background.radial_comoving_distance']
<dict 130500599962048> is referred to by <Cosmology 130500600156928>._workspace
<Cosmology 130500600156928> is referred to by <tuple 130500600588864>[5]

Traced<ShapedArray(float32[256])>with<DynamicJaxprTrace>
The error occurred while tracing the function _fn at /home/wassim/micromamba/envs/jax/lib/python3.10/site-packages/equinox/_eval_shape.py:31 for jit. 
<DynamicJaxprTracer 130500598573680> is referred to by <dict 130500600235072>['chi']
<dict 130500600235072> is referred to by <dict 130500599962048>['background.radial_comoving_distance']
<dict 130500599962048> is referred to by <Cosmology 130500600156928>._workspace
<Cosmology 130500600156928> is referred to by <tuple 130500600588864>[5]

My workaround is not good enough since cosmo is not passed as args .. knowing that diffrax guarentees that the adjoint of args is probably handled .. if it is passed around like the wrapper example I don't think the gradients are propagated correctly.
Also pre_computing is not very clean

tl;dr : I think that the jax_cosmo caching system can be more robust probably by using compiler directives such as jax.ensure_compile_time_eval

I will start making a PR when I get the chance.
In the mean time tell me what do you think

[tilmantroester]

I've just run into the same problem. Have you been able to reproduce the leak without diffrax?

[ASKabalan]

Hello @tilmantroester

Yeah it is related to the issue with the eval_shape that evaluates with a traced cosmo object .. so it is traced with no value and so the cache is written twice and thus we have leak

solution is to write a factory that creates your odeterm with arg = None

in otherwords

# BAD

def ode_with_chi(t, y, cosmo):
    distance = jc.background.radial_comoving_distance(cosmo, t).squeeze()
    return jnp.squeeze(distance) * 0.1 + y

ode_term = ODETerm(ode_with_chi)

# Good

def ode_wrapper(cosmo):
    def fn(t, y, args):
        # args here is None in diffeqsolve
        return ode_with_chi(t, y, cosmo)

    return fn

ode_term = ODETerm(ode_wrapper(cosmo))

the gradients still propagate correctly since cosmo is never (or at least not supposed to) change inside the kick and drift terms

[tilmantroester]

I found a similar workaround (in my case the solver and step size controller also depend on cosmo, so I had to find a different place to pre fill the cosmo._workspace).

Do you have an example with eval_shape that reproduces this leak? I tried but haven't been able to make an example that doesn't go through the diffrax machinery.

[ASKabalan]

It is difficult to reproduce this outside of diffrax
But last time I dove into this, I was able to avoid the error if commented out this line

https://github.com/patrick-kidger/diffrax/blob/90bf627f4dc55850d95864b0c212c3fdb3022e5e/diffrax/_integrate.py#L165

One more question?
Are you only using jax_cosmo or also jaxpm?