Using value_and_grad and jvp in same code snippet

[kopanicakova]

Hi all,

I am very new to JAX. My goal is to implement a Newton optimizer using JAX. So, I would like to use JAX to compute the gradient and evaluate the hessian-vector application of my function. I can perform both operations separately. However, if I use both features in the same code, I am getting the error "TypeError: 'DeviceArray' object is not callable". Can anybody point me to the problem? Below, you can find my current snippet:

def rosenbrock(x):
    return ((1-x[0])**2 + 100*(x[1]-x[0]**2)**2)

def hessian(f):
    return jax.jacfwd(jax.jacrev(f))

def hvp(f, primals, tangents):
  return jax.jvp(grad(f), primals, tangents)[1]    

x = jnp.zeros(2, dtype=jnp.float32)
v = jnp.ones(2, dtype=jnp.float32)

f, grad = value_and_grad(rosenbrock)(x)
HvP = hvp(rosenbrock, (x,), (v,))

[kopanicakova]

Figured it out - unfortunately, I have used name "grad" for output of value_and_grad function....