How to best code high order custom derivatives that require l'Hôpital's rule internally?

[flywithmath]

I am implementing methods using radial basis functions in RKHS context and want to use JAX autodiff capabilities.

For a vector z, a (strictly) positive definite matrix M and a scalar function phi: IR -> IR consider the functions

n(z) := Sqrt[<z, M z>]
k(z) := phi(n(z)).

For z != 0, we can compute the gradient

grad(k)(z) = phi'(n(z)) * (M z / n(z)).

Since typically: phi'(0) = 0 and phi''(0) != 0, we can use l'Hôpital (or Taylor expansion of phi' in zero) to define grad(k)(0) = 0 in a continuous fashion. I can implement this in JAX with custom_jvp.

However, I want/need to compute higher order derivatives in my method. The continuation using l'Hôpital's rule is need for any derivative in z=0 as the chain rule continues to produce 1/n(z) terms.

How would I best implement the high order derivatives in JAX and at the same time have them available for JAX' autodiff for further constructions?

An example for phi is Wendland's function(s) (https://num.math.uni-goettingen.de/picap/pdf/E392.pdf). There is also an other topic to be taken care for Wendland's function when differentiating outside the unit ball but this is not the topic of my question.

Also note, if phi(x) = psi(x^2), then the problem disappears if psi is smooth, as n(z)^2 = <z, M z> with gradient equal to 2 M z is also smooth. Composition works nicely in this case. This solves the problem for Gaussian's and inverse multi-quadric kernels say.

Any thoughts/hints for an implementation are highly appreciated.

[shoyer]

The short answer is that you implement your custom derivative rule in terms of a function with another custom derivative rule defined!

I can give more concrete advice if you can share a code snippet that isn't working as expected.

[jakevdp]

One example of this kind of approach can be found in the custom_jvp for the sinc function: https://github.com/google/jax/blob/be1cf46a49d260e979c68a78348ce7ce931e8151/jax/_src/numpy/ufuncs.py#L715-L740

[flywithmath]

Thanks for your food of thought. I get the idea. Great help!

[flywithmath]

Inspired by the proposal, I implemented a solution to compute 4th order derivatives of radial basis function kernels k(x,y) = phi(||x-y||) with ||z||^2 := <z, M z> for some positive definite matrix M. Doing all the algebra up to 4th order is quite cumbersome...

I thought I share the solution

import jax
import jax.numpy as jnp

@partial(jax.jit, static_argnums=(0,))
def eval_rbf(phi, m, x, y):
r"""Evaluate a radial basis function kernel k(x,y) := phi(||x-y||) with ||z||^2 := <z, m z>"""

def r2(x,y):
    xn2: Final = jnp.einsum('...ij,jk,...ik->...i', x, m, x)
    yn2: Final = jnp.einsum('...ij,jk,...ik->...i', y, m, y)
    x_dot_y: Final = jnp.einsum('...ik,kl,...jl->...ij', x, m, y)
    y_dot_x: Final = jnp.einsum('...ik,kl,...jl->...ij', y, m, x)
    
    return xn2[..., :, None] + yn2[..., None, :] - (x_dot_y + y_dot_x)

def r(x,y):
    return jnp.sqrt(r2(x, y) + 0*1j).real

r_xy = r(x,y)
near_zero = jnp.abs(r_xy) <= 1e-6

return jnp.where(near_zero, eval_rbf_on_diagonal_0(phi, m, x, y), phi(r_xy))

@partial(jax.custom_jvp, nondiff_argnums=(0,1))
def eval_rbf_on_diagonal_0(phi, m, x, y):
return phi(0.0)*jnp.ones(shape=(x.shape[0],y.shape[0]))

@partial(jax.custom_jvp, nondiff_argnums=(0,1))
def eval_rbf_on_diagonal_1(phi, m, x, y, h):
r"""Compute the 1st derivative at of k(x,y) at x=y denoted by dk(x,x)(h1) with
h1=(h1x, h1y) perturbation direction matrix of shape of x resp. y"""
return jnp.zeros(shape=(x.shape[0],y.shape[0]))

@partial(jax.custom_jvp, nondiff_argnums=(0,1))
def eval_rbf_on_diagonal_2(phi, m, x, y, h2, h1):
r"""Compute the 2nd derivative at of k(x,y) at x=y denoted by d^2k(x,x)(h2,h1)
h2=(h2x, h2y), h1=(h1x, h1y) with h1x, h2x, h1y, h2y perturbation direction matrices of shape of x resp. y"""
ht = jnp.array((h2, h1))

def q(a,b):
    return jnp.einsum('...k,kl,...l->...', a, m, b)

idx: Final = jnp.array([[0,0],
                        [0,1],
                        [1,0],
                        [1,1]])    
quad_part, _ = jax.lax.scan(lambda carry, i: (carry + (-1)**(i[0]+i[1])*q(ht[0,i[0]], ht[1,i[1]]), None),
                            jnp.zeros(shape=(x.shape[0], y.shape[0])), idx)        
phi_2_0 = jax.grad(jax.grad(phi))(0.0)

return phi_2_0 * quad_part

@partial(jax.custom_jvp, nondiff_argnums=(0,1))
def eval_rbf_on_diagonal_3(phi, m, x, y, h3, h2, h1):
r"""Compute the 2nd derivative at of k(x,y) at x=y denoted by d^3k(x,x)(h3,h2,h1)
h3 = (h3x, h3y), h2=(h2x, h2y), h1=(h1x, h1y) with h1x, h2x, h3x, h1y, h2y, h3y perturbation direction
of shape of x resp. y"""
return jnp.zeros(shape=(x.shape[0],y.shape[0]))

def eval_rbf_on_diagonal_4(phi, m, x, y, h4, h3, h2, h1):
r"""Compute the 2nd derivative at of k(x,y) at x=y denoted by d^4k(x,x)(h4,h3h2,h1)
h4 = (h4x, h4y), h3 = (h3x, h3y), h2=(h2x, h2y), h1=(h1x, h1y)
with h1x, h2x, h3x, h4x, h1y, h2y, h3y, h4y perturbation direction matrices of shape of x resp. y"""
def q(a,b):
return jnp.einsum('...k,kl,...l->...', a, m, b)

def q4(a,b,c,d):
    return q(a,d)*q(b,c) + q(b,d)*q(a,c) + q(a,b)*q(c,d)

idx: Final = jnp.asarray(((0, 0, 0, 0),
                          (0, 0, 0, 1),
                          (0, 0, 1, 0),
                          (0, 0, 1, 1),
                          (0, 1, 0, 0),
                          (0, 1, 0, 1),
                          (0, 1, 1, 0),
                          (0, 1, 1, 1),
                          (1, 0, 0, 0),
                          (1, 0, 0, 1),
                          (1, 0, 1, 0),
                          (1, 0, 1, 1),
                          (1, 1, 0, 0),
                          (1, 1, 0, 1),
                          (1, 1, 1, 0),
                          (1, 1, 1, 1)))

ht = jnp.array((h4,h3,h2,h1))

fourth_part, _ = jax.lax.scan(lambda carry, i: (carry + (-1)**(i[0]+i[1]+i[2]+i[3])*q4(ht[0,i[0]], ht[1,i[1]], ht[2,i[2]], ht[3,i[3]]), None),
                            jnp.zeros(shape=(x.shape[0], y.shape[0])), idx)   

phi_4_0 = jax.grad(jax.grad(jax.grad(jax.grad(phi))))(0.0)

return phi_4_0 * fourth_part / 3.0

@eval_rbf_on_diagonal_0.defjvp
def eval_rbf_on_diagonal_0_jvp(phi, m, primals, tangents):
x, y = primals
return eval_rbf_on_diagonal_0(phi, m, x, y), eval_rbf_on_diagonal_1(phi, m, x, y, tangents)

@eval_rbf_on_diagonal_1.defjvp
def eval_rbf_on_diagonal_1_jvp(phi, m, primals, tangents):
x, y, h1 = primals
h2x, h2y, _ = tangents
h2 = (h2x, h2y)

return eval_rbf_on_diagonal_1(phi, m, x, y, h1), eval_rbf_on_diagonal_2(phi, m, x, y, h2, h1)

@eval_rbf_on_diagonal_2.defjvp
def eval_rbf_on_diagonal_2_jvp(phi, m, primals, tangents):
x, y, h2, h1 = primals
h3x, h3y, _, _ = tangents
h3 = (h3x, h3y)

return eval_rbf_on_diagonal_2(phi, m, x, y, h2, h1), eval_rbf_on_diagonal_3(phi, m, x, y, h3, h2, h1)

@eval_rbf_on_diagonal_3.defjvp
def eval_rbf_on_diagonal_3_jvp(phi, m, primals, tangents):
x, y, h3, h2, h1 = primals
h4x, h4y, _, _, _ = tangents
h4 = (h4x, h4y)

return eval_rbf_on_diagonal_3(phi, m, x, y, h3, h2, h1), eval_rbf_on_diagonal_4(phi, m, x, y, h4, h3, h2, h1)