Speed of @tf.function and jax.jit

[pipme]

I am computing value and gradient of a function f using both tensorflow and jax. For tensorflow, I decorate the value and gradient computation with tf.function. The implementations of f in tensorflow and jax are basically equivalent (e.g. tf.linalg.cholesky -> jax.scipy.linalg.cholesky, tf.linalg.triangular_solve -> jax.scipy.linalg.solve_triangular). The table below shows the time for the first call on the value and gradient computation and the average time fo latter computations, tested on a P100 GPU.

	tf.function(jit_compile=False)	tf.function(jit_compile=True)	jax.value_and_grad(f)	jax.jit(jax.value_and_grad(f))
Frist call /s	2.559	8.213	8.208	3.909
Average after first call /s	0.0191	0.0153	0.0664	0.0158

I am a bit confused by two things.

1. jax.value_and_grad(f) takes more time for the first call than jax.jit(jax.value_and_grad(f)). Is it normal and why this happens? I would have presumed that jax.jit takes more time as it need to do both compilation and tracing.
2. tf.function(jit_compile=False) performs much better than jax.value_and_grad(f) and takes less time than jax.jit(jax.value_and_grad(f)) on the first call. So tf.function(jit_compile=False) could be more efficient compared to jax.jit for some cases, e.g. when we compute 10 times of f's value and gradients.

Also, jax.scipy.linalg.cholesky gets jitted no matter we call jax.jit or not, right?

Thanks.

[YouJiacheng]

Hmmm, IIUC tf.function doesn't compute gradient for you?

[pipme]

I use the following function to compute value and gradient, decorated with tf.function : )

@tf.function(jit_compile=True|False)
def _compute_loss_and_gradients(
    f,
    variables,
):
    with tf.GradientTape(watch_accessed_variables=False) as tape:
        tape.watch(variables)
        value = f()
    grads = tape.gradient(f, variables, unconnected_gradients=tf.UnconnectedGradients.ZERO)
    return value, grads

[hawkinsp]

In a sense, JAX has two tiers of compilation ("op by op", jit) where TF has three ("op by op", tf.function without jit, tf.function with jit). We have sometimes talked about adding a third tier to JAX between "op by op" and jit, but have never yet found a use case where someone needs that middle tier.

If you have a compelling use case where you need an intermediate point between the two, we'd love to hear it. Usually the right design point if you care about speed is "use jit".

By the way, even if you don't jit your code, many JAX library functions use jit internally. There's no real notion of an "op" in JAX distinct from jit: we have primitives (the operations mainly in jax.lax) and most other functions we offer are simply prebuilt jit functions that compose those primitives together. So for example jax.scipy.linalg.cholesky is indeed just a jit-decorated function https://github.com/google/jax/blob/b666f665ece2b91e7c07c65d27cf485a779638b6/jax/_src/scipy/linalg.py#L40

For "jax.value_and_grad(f) takes more time for the first call than jax.jit(jax.value_and_grad(f))", what that really says is that compiling N small computations took more time than one big one. I don't know why that is the case without digging into it, but I will note that one possible advantage of the small computations is that that you may be able to get more compilation cache reuse between similar but not identical computations if you they are smaller. But you'll almost certainly get the best performance out of one big jit once the compilation is cached.

[pipme]

Thanks for the super clear answer! In my case, the array input to f keeps changing shape since I am actively obtaining new data. Therefore, I need to do jax.jit(jax.value_and_grad(f)) each time new data arrives for l-bfgs optimization, which slows down a bit the algorithm. But for my concrete case the intermediate point between the two (like tf.function without jit) may not be so compelling, since the l-bfgs optimization often takes ~1000 calls of jax.jit(jax.value_and_grad(f)) and dominates the runtime compared to jit compilation. One compelling case I can think of is that the optimization only takes tens of calls. Then @tf.function without jit is significantly better than JAX jit, as shown in the table.