Unclear how jax.experimental.jet works

[Sceki]

Hello,

I am using jet, as we have a neat case where Jet Transport derivatives are very useful. However, I cannot really make sense of the current API even for very simple cases.

Is there any way I can extract the partials separately?

For instance, for the following function:

import jax
import jax.numpy as jnp
from jax.experimental.jet import jet

def f(x,y):
    return jnp.sin(x)*jnp.cos(y)

if I pass:

primals=(0.,0.)
series=[(1.,0.,0.,0.,0.,0.,0.), (1.,0.,0.,0.,0.,0.,0.)]

The output:

primals_out, series_out = jet(f, primals, series)

returns a series_out that is a list of jax arrays, as long as each tuple in series: how should I interpret that?

More generally, is there a way to obtain all the Taylor coefficients, evaluated at the primal point? E.g.:
$$\dfrac{\partial f}{\partial x}, \dfrac{\partial f}{\partial y},\dfrac{\partial^2 f}{\partial x^2} , \dfrac{\partial^2 f}{\partial y^2} , \dfrac{\partial^2f}{\partial x \partial y} , \dots$$

[Sceki]

Perhaps @mattjj could help here, as I have seen him chiming in some jet discussions :)