In JAX, why does using jvp with vmap yield a different Jacobi matrix than using jacfwd?

[tulvgengenr]

https://jax.readthedocs.io/en/latest/tutorials/advanced-autodiff.html#the-implementation-of-jax-jacfwd-and-jax-jacrev

In the documentation, I found a difference between the results obtained using jvp modified with vmap and directly using jacfwd.

I also verified this result in my experiment. But the strange thing is that I have the same results on the CPU (atol=1e-8), but there is a difference in the results on the GPU (V100) (atol=1e-5).

Although the documentation explains that their results are different, I still don't quite understand why.

Also, in my scenario, I need to compute the entire Jacobi matrix, the diagonal entries of the Hessian matrix, and using the forward differentiation mode. For the Hessian matrix, should I use jvp over jvp to compute it, or should I use jacfwd over jacfwd and then pick the index to compute it?

Thanks!