[jax-tutorial 101] how to jit

[JiahaoYao]

In the tutorial, when to jit says that

Generally, you want to jit the largest possible chunk of your computation; ideally, the entire update step. This gives the compiler maximum freedom to optimize.

Here are the codes I have at a high level, I wonder how to correctly jit the function!

The function is this

def sovler(x): 
    for i in range(1000): 
        x = update1(x, ...)
        x = update2(x, ...)

    return x

solution = [solver(x) for x in x_list] 

---

approach I

@jit 
def sovler(x): 
    for i in range(1000): 
        x = update1(x, ...)
        x = update2(x, ...)

    return x

solution = [solver(x) for x in x_list] 

do I just run this? (because this function is also called multiple time)

approach II

or I also do this

@jit 
def update1(x):
    ... 

@jit 
def update2(x): 
    ... 

@jit 
def sovler(x): 
    for i in range(1000): 
        x = update1(x, ...)
        x = update2(x, ...)

    return x

solution = [solver(x) for x in x_list] 

This is because the update1 and update2 also need to run multiple times.

---

I wonder

1. whether there is any difference between the two approaches, if not, does that mean the jax will automatically jit the update1 and update2 to optimize this big block of code.
2. a jit function inside a jit function will slower the speed or not?

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

[jakevdp]

Thanks for the question! I don't think this is well-covered in the tutorials, but from the user perspective there is essentially no difference between your approach 1 and approach 2: everything that is computed within a jitted function is JIT-compiled; if parts of that happen to be separately jitted, it makes no difference.

One comment, though: JIT compilation will unroll Python loops, which means that as written there will be 1000 copies of your operations serialized and passed to the compiler in both approaches, so either one will likely lead to very slow compilation times. One way to work around this is through Structured control flow operations.

[JiahaoYao]

Hi @jakevdp , thank you again for your awesome reply. Your comments on using the control flow actually are wonderful suggestions! Thanks!