A function with default arguments was much faster in JAX. How did it happen? How does JAX trace its gradient too?

[ToshiyukiBandai]

Hi all,

I experienced a performance gain if I used default arguments in a function. This is related to the discussion on stack overflow (https://shorturl.at/ahmRS). In the example below, I compared the performance of a function with and without default arguments, and the former was 94.3 us and the latter was 220 us on my computer. So, I checked their jaxpr and noticed that XLA optimized alpha**2*beta when they are given as default arguments. I suspected that I can't get the gradient of the function with default arguments, but it worked as well. So, a good practice is to use as many default arguments as possible?

import timeit
import jax
from jax import lax, grad
import jax.numpy as jnp
import timeit
from jax import make_jaxpr

x = jnp.linspace(0, 1, 100)

def f(x, alpha, beta):
    x2 = alpha**2*beta*x**2
    return x2.sum()

def f2(x, alpha = 2.0, beta = 3.0):
    x2 = alpha**2*beta*x**2
    return x2.sum()

alpha = 2.0
beta = 3.0

jitted_f = jit(f)
jitted_f2 = jit(f2)

# warm up and print jaxpr
print(make_jaxpr(jitted_f)(x, alpha, beta))

print(make_jaxpr(jitted_f2)(x))

%timeit jax.block_until_ready(jitted_f(x, alpha, beta))

%timeit jax.block_until_ready(jitted_f2(x))

# both grad works!
# grad(f, argnums=(0, 1, 2))(x, alpha, beta)
# grad(f2, argnums=(0, 1, 2))(x, alpha, beta)

[mattjj]

Thanks for the question!

I would guess that jitted_f2(x) can be faster because the compiler gets to specialize on the values of alpha and beta as constants rather than taking them as arguments at execution time, whereas jitted_f(x, alpha, beta) basically means "ask the compiler to make an executable taking alpha and beta as arguments only to be specified when it's called".

We could check by looking at the optimized program, as described here.

What do you think?

[ToshiyukiBandai]

Hi Matthew,

Thank you for answering my question! I read through the suggested document and checked the optimized program (lowered HLO). As you said, the optimized program for jitted_f2(x) was smaller. But, the optimized program for jitted_f2(x, alpha2, beta2) was the same as the one for jitted_f(x), which is reasonable; JAX adjusted to the new arguments. I guess I should use default arguments whenever possible.