Higher Order and Mixed Derivatives of Multivariate Functions (jax.experimental.jet)

[jmsdsouzaPhD]

Hello,

I'm Josiel de Souza, postdoc in Physics working with gravitational wave data analysis.

I'm interested to compute high order derivatives of functions.

I heard that using jax.experimental.jet is more efficient than using jax.grad recursively (the later is very slow), and, indeed, I saw that it works very well for 1-dimensional functions f(x). However I don't know to to use jax.experimental.jet to compute mixed derivatives $\frac{\partial^2f}{\partial x \partial y}$ as well as individual derivatives. I mean, I'd like to compute $\frac{\partial f}{\partial x}$, $\frac{\partial f}{\partial y}$, $\frac{\partial^2f}{\partial x\partial y}$, $\frac{\partial^2 f}{\partial x^2}$, $\frac{\partial^2f}{\partial y^2}$ and so on using jax.experimental.jet. I don't understand how the input series work in this case.

Could anyone help me with a practical example, for instance, for the function $sin^2x/y$?