Partial versus static_argnums for passing functions into jitted function

[KeAWang]

Suppose we want to be able to pass a closure function (possibly jitted) into another jitted function. I think generally, we do one of two things shown below by f1 and f2.

import jax
from functools import partial

@partial(jax.jit, static_argnums=0)
def f1(g, x):
    return g(x)

@jax.jit
def f2(g, x):
    return g(x)

y = 0
# @jax.jit  #maybe we want to jit this too
def g(x):
    return x + y

x = 0

print(f1(g, x))
print(f2(jax.tree_util.Partial, x))

Which one is preferable? And what are the tradeoffs for each?

[jakevdp]

Did you mean for the second call to be this?

print(f2(jax.tree_util.Partial(g), x))

Regardless, the general pattern that people use is either to use static_argnums/static_argnames or use a closure. Side-by-side, they'd look like this:

def f(g, x):
    return g(x)

ans1 = jax.jit(f, static_argnames=['g'])(g, x)
ans2 = jax.jit(partial(f, g))(x)

Both are idiomatic, and don't really have any tradeoffs; you can use whichever fits the context of your code. If you want f to be a stand-alone JIT-compiled function (i.e. not leave JITting up to the caller) then static_argnames is probably best.

Using tree_util.Partial also works, the only tradeoff might be that it is more indirect and may be confusing to someone coming across it in your code.

[KeAWang]

As a followup, would jax.jit(Partial(g)) and jax.jit(g) share the same jit cache? I think jit duplicates closed variables (y in this case), so I'd like to prevent duplicating closed variables in memory.

[jakevdp]

No, I don't believe they'd share the same cache because Partial(g) is a pytree and the pytree hash is taken into account when computing the JIT cache key.

[KeAWang]

Is there a way to see what all the JIT cache key value pairs are?

[jakevdp]

No, you can't directly examine cache keys as far as I know, but you can introspect the cache size if that's helpful:

import jax
from jax.tree_util import Partial

g = lambda x: x
jit_g = jax.jit(g)
jit_partial_g = jax.jit(Partial(g))

print(jit_g._cache_size())
# 0

jit_g(1)
print(jit_g._cache_size())
# 1

# Jit-compiling the same function again shares the original cache
print(jax.jit(g)._cache_size())
# 1

# but Partial(g) != g, so it does not share the cache
print(jit_partial_g._cache_size())
# 0

[KeAWang]

This is great! Thank you so much!