Integration in jax

[jschiavon]

Hi guys,
first of all thanks a lot for your work, I am loving this library!

For my work I need to optimize (i.e. diffferentiate) a function which has an integral inside (quite simple, as it is one-dimensional) that depends on one of the variable I want to differentiate on.
I have done a search around and I have not found any package dealing with jax-compatible integration methods, am I right?

If this is the case, I might have to implement a couple of methods myself (probably simple quadrature is enough for my needs), but I hope not to!

[PhilipVinc]

Buongiornissimo.

I achieved something like that in the past by using jax.custom_vjp/jax.custom_jvp (depending on whever you need fwd/bwd differentiation). See here for some examples.

Essentially you can define a function like

@partial(jax.custom_vjp, nondiff_argnums=0)
def my_fancy_fun(integration_mode, integral_params):
    # this can be not differentiable
    return integrate(my_integrand, *integral_params, mode=integration_mode)   

def my_fancy_fun_fwd_pass(integration_mode, integral_params):
    # this can be not differentiable
    fwd_res = integrate(my_integrand, *integral_params) 
    fwd_data = None
    return fwd_res, fwd_data

def my_fancy_fun_fwd_pass(integration_mode, fwd_data, fwd_res):
    # this can also be not differentiable
   def my_grad(vector)...
   
   return my_grad

my_fancy_fun.defvjp(my_fancy_fun_fwd_pass, my_fancy_fun_fwd_pass)

in the body of the forward and backward pass you can use anything even if it is not differentiable, but it must be jax compatible.

Now, if you want to call back into standard scipy, my usual trick is to use

def integrate(f, *pars, mode="something"):
   # f and mode must be hashable constants
   from jax.experimental import host_callback as hcb
   return hcb.call(partial(_integrate, f, mode=mode), pars, result_shape=jax.ShapeDtypeStruct((), jnp.float64)

def _integrate(f, *pars, mode):
   return scipy.integrate(f, pars, mode)

This calls back into python. You must be able to tell XLA the result_shape (so shape and dtype) of the output.
If you do this both for the forward and backward pass, and combine this with custom_vjp to get autodiff working, you don't need to rewrite the integration code.

[jschiavon]

Thank you a lot!!!

I always managed to avoid actually digging into custum_vjp/jvp but it seems I cannot avoid it anymore!

(e buongiornissimo a te! 😄 )

[pharringtonp19]

@PhilipVinc Do you know of a "vmappable" method? I get the following error when I wrap my function with vmap

NotImplementedError: batching rules are implemented only for id_tap, not for call.

I am somewhat surprised that call() works for pmap but not vmap. -- see here: "The call() API can be used inside a jit or pmap computation or inside cond/scan/while control flow."

[PhilipVinc]

The reason that it works for pmap is that the extra axis is never passed on to the function call, i think (though jax people may give you better explanations).

A custom_vmap for a call function would be handy.
If you are familiar with jax internals, you could implement the same system used for jax.custom_vjp to make a custom_vmap function, but you'll need to define a new primitive and it's vmap tracer.

See here for inspiration https://jax.readthedocs.io/en/latest/_modules/jax/_src/custom_derivatives.html#custom_vjp