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$?

[DavidLanders95]

Pretty sure it is explained here at least a little bit better: #25700

"Each element of the outer list corresponds to the taylor series of one of the elements of the primal list, and the length of the element dictates the degree of the truncated Taylor polynomial. For example: primals = (x, y), series = ((dx, ddx), (dy, ddy))."

and here is an example of someone calculating partial derivatives also: #5152

Although, I too don't quite understand how they are doing it. Did you ever figure this out?

[DavidLanders95]

I had another look at this, and with some copilot help was able to get an example that works I think.

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

jax.config.update("jax_enable_x64", True)

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

def f_sym(x, y):
    return sp.sin(x)**2 / y

def get_first_order_mixed_partials(func, variables):
    mixed_partials = {}
    for i in range(len(variables)):
        for j in range(i+1, len(variables)):
            # Compute the mixed partial derivative with respect to variables[i] and variables[j]
            d_expr = sp.diff(func, variables[i], variables[j])
            mixed_partials[(variables[i], variables[j])] = sp.simplify(d_expr)
    return mixed_partials

#Calculate symbolic derivative with sympy
x, y = sp.symbols('x y')
mixed_derivs = get_first_order_mixed_partials(f_sym(x, y), [x, y])
print("Mixed partial derivative (d²f/dxdy):", mixed_derivs[(x, y)])

primals = jnp.array([5.0, 1.0])
df_dxdy_func = sp.lambdify([x, y], mixed_derivs[(x, y)], 'numpy')
df_dxdy_sympy = df_dxdy_func(*primals)

#Calculate symbolic derivative with jax
series = [
    [jnp.array([1.0, 1.0]), 0.0],   # Perturb along (1, 1)
    [jnp.array([1.0, -1.0]), 0.0]   # Perturb along (1, -1)
]
primal, jet_coeffs = jet(f, primals, series)

d_series1 = jet_coeffs[1][0]  # from series for direction (1, 1)
d_series2 = jet_coeffs[1][1]  # from series for direction (1, -1)

# By Taylor expansion:
#   f(x + eps*d) = f(x) + eps*(grad f dot d) + (eps^2)/2*(d^T Hessian d) + ...
# For the (1,1) direction, the quadratic term is 0.5*(fₓₓ + 2·fₓᵧ + fᵧᵧ)
# and for (1,-1) it is 0.5*(fₓₓ - 2·fₓᵧ + fᵧᵧ).
# Subtracting these yields:
#   0.5*( (fₓₓ+2·fₓᵧ+fᵧᵧ) - (fₓₓ-2·fₓᵧ+fᵧᵧ) ) = 2 · fₓᵧ.
# Since jet returns the (unscaled) Taylor series coefficients for the quadratic term,
# we combine them to solve for fₓᵧ.
# In our example, the mixed partial derivative is:
df_dxdy_jax = (d_series1 - d_series2) / 4

print("The mixed partial derivative via sympy: (d²f/dxdy) =", df_dxdy_sympy)
print("The mixed partial derivative via jax: (d²f/dxdy) =", df_dxdy_jax)

Mixed partial derivative (d²f/dxdy): -sin(2*x)/y**2
The mixed partial derivative via sympy: (d²f/dxdy) = 0.5440211108893698
The mixed partial derivative via jax: (d²f/dxdy) = 0.5440211108893697

The tricky part (if my understanding is correct of jet) is that jet is returning coefficients of an unscaled Taylor expansion, not just giving you back a gradient like jax.grad or jax.jacobian, and that you need to extract the partial derivatives from there. Maybe there is a cleaner method to extract partial derivatives, but for now this is all I can come up with. For this example I have just found $\frac{\delta f^2}{\delta x\delta y}$, and the general case will take a bit more work. I think more examples in the documentation would be very helpful for this for certain.

[DavidLanders95]

Nevermind my implementation, a much better answer is given here:

#5152 (comment)