y1xiaoc/fwdlap

Forward mode laplacian implemented in JAX tracer

fwdlap's Issues

API that takes only a function

Currently the lap function is somehow more like the jvp function, which takes not only the function, but also a list of primals, jacobians and laplacians.

One may expect the laplace transformation to take only a function and return a transformed function that shares the same input space of the original function. For example, like the jax.jacfwd.

Is there occasions that we want to manually specify primals, jacobians and laplacians? Or is it due to technical difficulty?

apply `lap` to vmap'ed function with custom jvp will fail

Currently lap support normal custom jvp calls, but they cannot be vmap'ed. For example, the following code would throw an error

with jax.softmax_custom_jvp(True):
    o0, o1, o2 = jax.jit(lap, static_argnums=0)(jax.vmap(jax.nn.softmax), (x[None],), 
                                                (jnp.eye(x.size)[:, None],), (jnp.zeros_like(x)[None],))

Instead, if possible, move vmap outside of lap should work.
Or, one may need to manually separate the function with custom jvp and do the forward rule for that function by hand.

This is because the current implementation handles custom_jvp calls in a very hacky way. To avoid an error from post_process_custom_jvp_call in ad.jvp, I have to remove the process_env_traces_fwd transformation in the wrapped custom jvp function. (Search unwrap in the code to see it.) I do not know how to fix this for now. Any advice or help is greatly appreciated!