Best practices for compilation OOM

[OhadRubin]

Hey,
What are some best for writing jax code that doesn't OOM during compilation?
If understand correctly it uses scans, but I haven't found any specific information related to this.
Thanks!

[jakevdp]

Hard to say without more information, but one thing to keep in mind is that JAX compilation unrolls python control flow. So something like

@jit
def f(x):
  for i in range(1000000):
    x += 1
  return x

would result in 1,000,000 add statements in the IR. So one best practice is don't use large Python for loops (see Sharp Bits: Control Flow for more on this).

[OhadRubin]

Ok, awesome.
Regarding more information: I am experiencing OOM in compilation after increasing the batch size, my implementation is using jnp.vectorize to handle the multiple dimensions, is that not advisable?

[jakevdp]

If jnp.vectorize is causing issues, are you certain the OOM is at compilation time and not at runtime? It's possible with vectorize to inadvertently create very large arrays given small inputs that might OOM at runtime.

[OhadRubin]

It happens during jitting, so I assumed that means compilation.
Using vmap would reduce the memory consumption compared to jnp.vectorize?

[jakevdp]

No, vectorize is just a wrapper around vmap (or multiple vmaps)

If you could edit your question with a minimal reproducible example of the issue, it may be possible to give more pointed advice.