Add support for jacrev, jacfwd, hessian, vmap

[martinjm97]

Extending the tests in

Enzyme-JAX/test/test.py

Line 68 in d9e2ae0

print(grads)

with

> x = jax.jacrev(add_one)(jnp.array([1., 2., 3.]), jnp.array([1., 2., 3.]))
NotImplementedError: Batching rule for 'enzyme_rev' not implemented

> x = jax.jacfwd(add_one)(jnp.array([1., 2., 3.]), jnp.array([1., 2., 3.]))
NotImplementedError: Batching rule for 'enzyme_fwd' not implemented

> x = jax.hessian(add_one)(jnp.array([1., 2., 3.]), jnp.array([1., 2., 3.]))
NotImplementedError: Differentiation rule for 'enzyme_aug' not implemented

> x = jax.jit(jax.vmap(lambda x: add_one(x, jnp.array([1., 2., 3.]))))(jnp.array([jnp.array([1., 2., 3.])]*5))
NotImplementedError: Batching rule for 'enzyme_primal' not implemented

[croci]

Are jax.grad and jax.value_and_grad already supported?

[martinjm97]

Hi @croci,

I believe so. I just tested it by adding:

x = jax.grad(add_one)(.1, .1)
x = jax.value_and_grad(add_one)(1., 2.)

to the end of the test.py file.

--Jesse

[croci]

@martinjm97 ok great! I couldn't get it to work on the pip version, but that is v0.0.4 (and I confess I didn't try too hard). Thanks a lot!