JAX does jitting functions separately change performances?

user_8737016 0
Viewed 136

I am learning to use JAX and I have some doubts about the use of jit and vmap that I couldn't solve by reading the docs.

  1. Does it make a difference to jit several functions separately and then jit the function that uses them? For example, if I have the functions foo() and bar() and a function

    @jax.jit 
def fooBar(x):
    return foo(x) + bar(x)

    Is there any difference if foo() and bar() are already jitted?

  2. Should I jit a function after I vmap it? In the example above, should I do jax.jit(jax.vmal(fooBar)) or just jax.vmap(fooBar)?

jax jit
1 Answers

When it comes to performance of code execution, there is no difference between jitting functions separately and jitting once at the outer function (functionally there is one subtle difference: jit-compiling the inner function will wrap the contents in an xla_call primitive, but this makes little to no difference for the final compilation & execution).

When using vmap on the other hand, there is no implicit compilation. vmap(f) will be executed in eager mode, while jit(vmap(f)) will be just-in-time compiled and generally result in faster execution.

user_29378310
Related