human
I am working on the file "Fall2020/03-MatrixCalculus/jax_autodiff_tutorial.ipynb".
The first task is:
Hessian-by-vector product
Can you write Python code for it?

Write me a the Python code you requested?
assistant
#Here is the Python code you requested.

def hvp(f, x, z, *args):
    def g(x):
        return f(x, *args)
    return jax.jvp(jax.grad(g), (x,), (z,))[1]

n = 3000
x = jax.random.normal(jax.random.PRNGKey(0), (n, ))
A = jax.random.normal(jax.random.PRNGKey(0), (n, n))
b = jax.random.normal(jax.random.PRNGKey(0), (n, ))
z = jax.random.normal(jax.random.PRNGKey(0), (n, ))

print("Check correctness", jnp.linalg.norm(2 * A.T @ (A @ z) - hvp(f, x, z, A, b)))
print("Time for hvp by hands")
%timeit (2 * A.T @ (A @ z)).block_until_ready()
print("Time for hvp via jvp, NO jit")
%timeit hvp(f, x, z, A, b).block_until_ready()
print("Time for hvp via jvp, WITH jit")
%timeit jax.jit(hvp, static_argnums=0)(f, x, z, A, b).block_until_ready()