OOM issue fori_loop

[JasonKt]

I'm new to jax, and I foolishly wrote a statistical function that uses a for loop with a nested fori_loop that swallows lots of memory on each iteration. Am rather stuck trying to work out a way to avoid out of memory--hopefully someone can steer me in the right direction. The function is, roughly as follows (note that probs is a 2D array):

def log_prob(self, probs):
        logp = 0.
        for i in range(len(probs)):
            p = probs[i]
            initialValPp = function spitting out a 1D array based on a transformed p
            out=lax.fori_loop(1, self.max_len_p, self.f_partial, initialValPp)
            pval = out[self.max_len_p + self.y[i]]
            ln_pval = jnp.log(pval)
            logp += ln_pval
        return logp

Each time the fori_loop executes, my memory declines by about 220MB. The initialValPp array is only about 88KB, though the fori_loop executes thousands of times. Am guessing that each of the executions somehow adds to the memory being consumed.

I only need one value from 'out'. So I tried extracting a float from 'out'. That solves the memory problem when the function isn't jitted, but later it needs to be jitted and then I get a ConcretizationTypeError as a direct result of extracting the float.

I also tried a mock up version of the function that gets rid of the 'for' loop and replaces it with a fori_loop. That didn't help with the memory problem.

Any suggestions as to what I can try are welcome!

[jakevdp]

Hi - it's difficult to guess what's happening here without a runnable example. Can you create a complete minimal reproduction of the issue, with simplified versions of the functions that you're calling?

[JasonKt]

Thanks for your reply jakevdp! I took the code out of its NumPyro context. When I do, the memory problem disappears. I went to some lengths to reproduce my inputs into the problem using numpyro.sample to generate initial parameter estimates. Still, no problem.

I'm guessing that there's something that NumPyro is doing behind the scenes that results in the big memory usage. Maybe it's the auto-differentiation. The fori_loop convolves, thousands of times, the initial parameters, so if it's keeping track of info for, say, backward differentiation, that may explain big memory use. Just hoping I can find some way to turn this off.

[JasonKt]

Setting forward_mode_differentiation=True in the NUTS call fixes the memory issue. Thanks again!