Is it safe to use dict.pop in jit, vmap, and other jax transformations?

[HeavyCrab]

jit cannot correctly handle side-effects, but the following code seems to work correctly:

from jax import make_jaxpr

def f(xx):
    print('tracing')
    xx.pop('aa')
    return xx

xx = {'aa': 1, 'bb': 2}
print(make_jaxpr(f)(xx))
print(xx)   # xx is not modified

print(jax.jit(f)({'aa':1, 'bb':2}))
print(jax.jit(f)({'aa':1, 'bb':3}))
print(jax.jit(f)({'aa':1, 'bb':2, 'cc':3}))

output:

tracing
{ lambda ; a:i32[] b:i32[]. let  in (b,) }
{'aa': 1, 'bb': 2}
{'bb': Array(2, dtype=int32, weak_type=True)}
{'bb': Array(3, dtype=int32, weak_type=True)}
tracing
{'bb': Array(2, dtype=int32, weak_type=True), 'cc': Array(3, dtype=int32, weak_type=True)}

So is it safe to use dict.pop in a jitted function? Or should we use FrozenDict instead?

[HeavyCrab]

In addition, I don't understand why an inplace function like xx.pop('aa') can be compiled by JIT without issues. Shouldn't it report some error like Tracer Leak?

[jakevdp]

Whether or not it's safe to use dict.pop within jax.jit depends on the context.

In the snippet you shared, it is safe: explicit dict arguments to jit-compiled functions will be passed through jax.tree_util flattening, and so the dict you are operating on within the function is a copy of the dict passed to the function, and mutating it will not affect the dict in the outer scope.

However, if you remove jax.jit or execute your function in a jax.disable_jit context, mutating the dict will mutate the dict in the outer scope. So that's something to be aware of if you're using this kind of pattern.