question about the vjp and grad

[yiminghwang]

Hi, there

I am a beginner in using Jax, I follow the tutorial https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html#vjps-in-jax-code, and have a question about vjp and grad,

 def func(params):
      return params[0] ** 2 + 2 *  params[1]

 params = jnp.asarray([1.,2.])
 g1 = jax.grad(func)(params)

 y, f_vjp = jax.vjp(func, params)
 g2 = f_vjp(jnp.ones(2,dtype=jnp.dtype(y))[0])

the output of g1 and g2 might be the same.
my question is there must be some reason why Jax do not use the following or similar way , i.e. g2 = f_vjp(jnp.ones(2,dtype=jnp.dtype(y))[0]) to calculate the gradient. Does anyone explain it? Many thanks

[soraros]

grad is implemented using value_and_grad, which is in turn implemented in terms of vjp.

[yiminghwang]

get it, thank you!