Memory consumption with jitted function that closes over large jnp.array #1222

[avital]

(cross-posted on google/flax#1222)

Just wanted to showcase a common gotcha (originally reported by @lucasb-eyer and answered by @hawkinsp. This should be documented somewhere else probably, tooling that @jheek is working can help with an explicit warning to users)

If you jit a function that closes over a jnp.array, then that array is compiled into the jitted function, and the Python reference is lost. This implies two things:

1. Changes to the variable will not be reflected in subsequent calls to the jitted function
2. Memory will be duplicated: both in the compiled program, and the original jnp.array.

In general, the workaround is to take in any arrays that are either large or may change as arguments to the jitted function, rather than closing over them.

For example, this small code snippet:

x = jnp.ones((1, 2))

@jax.jit
def f():
  return x

print(f())
print(x)

print("---")
x = x * 2
print(f())
print(x)

will print:

[[1. 1.]]
[[1. 1.]]
---
[[1. 1.]]
[[2. 2.]]