[jax-tutorial 101] how to have function output a jitted function

[JiahaoYao]

In the tutorial, when to jit says that

Avoid calling jax.jit inside loops. Doing that effectively creates a new f at each call, which will get compiled each time instead of reusing the same cached function:

    def g_inner_jitted_poorly(x, n):
    i = 0
    while i < n:
      # Don't do this!
      i = jax.jit(unjitted_loop_body)(i)
    return x + i

---

this is quite helpful, but I am a little confused about how to jit the following?

I want a function to output a function in order to pass some parameter?

def make_compute_fn(param1, param2): 
    ...
    def compute(param1, param2): 
        ...

    return compute

In order to have a jitted function, is this appropriate to do this?

def make_compute_fn(param1, param2): 
    ...
    @jit
    def compute(param1, param2): 
        ...

    return compute

Again, thanks for posting the Jax tutorials, really appreciate that!

[jakevdp]

Your last code snippet, returning the jitted function, would be the recommended way to do this.

The example you quote from the tutorial is about jitting within a loop, which leads to repeated JIT-compilation overhead. If you're creating a jitted function once within a factory function, that concern about repeated compilation overhead does not apply.

[JiahaoYao]

Hi @jakevdp , thank you so much for your quick comments!
One quick question, does this factory function amount to the following static_argnums?

    @jit(static_argnums=(0,1))
    def compute(param1, param2, x):
          something.... 

[jakevdp]

No, arguments to a jitted function will not be static unless they are explicitly marked as static. In your original example, param1 and param2 will not be static, they will still be traced.

I worry that this question is confusing two topics, though: leaving JAX and JIT out of the picture, if you write something like this:

def f(param1, param2):
  def g(param1, param2):
    return param1 + param2
  return g

The param1 and param2 passed to g have nothing to do with the param1 and param2 passed to f, despite having the same name.

[JiahaoYao]

Great! Your comments are exactly my previous confusion. Thanks a lot again!