Open jan-david-fischbach opened 1 year ago
Hum... Jax is actually build to return the result of the function AND the tangents using custom_jvp. The only way to do it will be to use a 'hack'. The simplest will be to define two different functions depending which mode to use (either 'value_and_grad' or 'grad'). But it is really dirty...
I seem to run into the problem, that if I use
jax.value_and_grad
I do not actually go through the forward pass of the generator. I think this is related to these lines: https://github.com/flaport/inverse_design/blob/86486d609d447a5247e5b8fe97f6e24ebdd376e3/inverse_design/conditional_generator.py#L188-L190I believe one could solve that by passing the primals through the generator in the customjvp. This however leads to unneccesary computational cost if I only want the gradients i.e. use the
jax.grad
function. Any Idea how to get it to work efficiently in both cases?