How to compute Jacobian of outputs w.r.t. inputs about objax HOT 3 CLOSED

1. Compute Jacobian.

You can reimplement object-oriented version of jacfwd or jacrev similarly how objax.GradValues is implemented.

Another option is to use combination of objax.Vectorize and objax.GradValues, in other words vectorize computation of gradient. DPSGD gradient module to some extend does it:

2. Why not mixing JAX and Objax

You can not use functional jax transformations with Objax (like jax.vmap, jax.pmap, jax.grad, etc...). In other words those transformations which takes function and returns a new function.
You can use other jax operations (for example all stuff from jax.numpy.*) safely with Objax.

All JAX primitives are stateless and pure functional (i.e. don't have and don't assume side-effects). Objax provides wrappers for JAX primitives to simplify state management and make is more natural for machine learning applications.

So if you try to mix JAX functional tranformations with Objax primitives it will break the state management and either code won't work at all or will work incorrectly.

3. Difference between JAX and Objax primitives

As I mentioned above, Objax provides wrappers which simplify state management.
So for example objax.Vectorize is a wrapper over jax.vmap which does the state management and enables usage of stateful Objax primitives with stateless JAX.

from objax.

I see, thank you for your explanation!

from objax.

1. Compute Jacobian.

You can reimplement object-oriented version of jacfwd or jacrev similarly how objax.GradValues is implemented.

It is hard to implement object-oriented version of jacfwd or jacrev, because jacfwd or jacrev do not support return auxiliary data. Is there solution? Thanks!

from objax.