Efficient implementation of Stratonovich drift correction

[ylefay]

Hello, thanks a lot for your work on this library!

I am trying to implement this function
$$b \mapsto \bigg(\sum_{j=1}^{n} \sum_{k=1}^{m} {b^{(j,k)}(x)\frac{\partial b^{(i,k)}(x)}{\partial x_j}}\bigg)_{i}, $$
which corresponds to the Ito to Stratonovich drift correction term for SDEs. Typically, $b$ is a function from $\mathbb{R}^n$ to $\mathbb{R}^{n\times m}$.

I wrote that

def drift_correction(diffusion, x):
    return jnp.einsum('jk,ikj->i', diffusion(x), jax.jacfwd(diffusion)(x))

and wanted to know if it was possible to implement it using jax.jvp, jax.vjp or jax.linearize.

WDYT?
Thanks.

Yvann

[anh-tong]

Hi, I think you approach is fair. As you mentioned, we can use jax.jvp to avoid the need for storing the matrix jax.jacfwd(diffusion)(x), which would otherwise require space of order $O(n \times m)$. Below is an example for you case (see this in Diffrax)

def drift_correction(diffusion, x):
    b = diffusion(x)
    _, corrected = jax.jvp(diffusion, (x, ), (b, ))
    return corrected

If you're interested in SDE, you may want to check out Diffrax which is designed for solving differential equations in JAX.