If $J$ is the Jacobian of some $f$ at some $x$, then we compute the jacobian using vmap on jvp or vjp, then we compute the Gramian by computing $JJ^\top$.
What we could do is to compose vjp and jvp and use vmap on that to compute the Gramian directly without conputing the full Jacobian. This should be much faster.
If$J$ is the Jacobian of some $f$ at some $x$ , then we compute the jacobian using $JJ^\top$ .
vmaponjvporvjp, then we compute the Gramian by computingWhat we could do is to compose
vjpandjvpand usevmapon that to compute the Gramian directly without conputing the full Jacobian. This should be much faster.