def f(x):
    return x ** 2

# grad returns a NEW FUNCTION that computes the derivative
df_dx = jax.grad(f)

print(f(3.0))      # 9.0
print(df_dx(3.0))  # 6.0 (derivative of x² is 2x, and 2*3=6)

        __
        __