Does jax.value_and_grad imply jit?

[ayaka14732]

From my understanding, jax.value_and_grad may imply jit because it needs to somehow compile the function.

[itk22]

Hi. As far as I understand, this is not the case. To be able to calculate the value and gradient of a function, all you need is a 'recipe' for differentiation, which is either inferred from the computational graph or defined with a custom_vjp. I attached an example of a function which cannot be jitted because it uses calls to original numpy, but can be differentiated because a custom_vjp is defined:

import jax
import numpy as np
import jax.numpy as jnp
from jax import custom_vjp

@custom_vjp
def f(x):
    # The function deliberately uses numpy functions which are not jittable
    return -(x[0] + np.sin(x[0])) * np.exp(-x[0]**2.0)


def f_fwd(x):
    return f(x), (x, )


def f_bwd(res, g):
    x = res[0]
    grad_x = np.exp(-x[0]**2) * (2 * x[0]**2 + 2 * x[0] * np.sin(x[0]) - np.cos(x[0]) - 1)
    return (np.array([grad_x * g]),)


f.defvjp(f_fwd, f_bwd)

print_errors = False  # Flag for printing the errors

# Check if the custom function is jittable
try:
    f_jitted = jax.jit(f)
    f_jitted(np.array([1.0]))
    print("Function is jittable")
except Exception as e:
    print("Function is not jittable")
    if print_errors:
        print("Error: ", e)

# Check if the custom function is differentiable
try:
    f_grad = jax.grad(f)
    f_grad(np.array([1.0]))
    print("Function is differentiable")
except Exception as e:
    print("Function is not differentiable")
    if print_errors:
        print("Error: ", e)

I hope this illustrates the point that you don't need to be able to jit a function to find its gradients. An important point to mention here is that the function will be compiled no matter if you use JIT or not. However, If you do use JIT (just in time compilation), you can benefit from the XLA optimizations.

[ayaka14732]

Thank you for your detailed explanation!

[jakevdp]

There's a more concise way to see that grad does not imply jit – grad uses concrete tracers rather than dynamic tracers, and can transform functions that are not jit compatible! For example:

import jax

def relu(x):
  return 0.0 if x < 0 else x

print(jax.grad(relu)(1.0))
# 1.0

print(jax.jit(relu)(1.0))
# ConcretizationTypeError: Abstract tracer value encountered...