Combining custom jvps and custom vjps for higher order derivatives

[Binbose]

Hey, I have a function f(r) that takes a 3-d vector r as input and internally, the function calculates the norm of the vector. I want to calculate the n-th derivative of f with respect to the norm |r| for very small norms |r| = 1e-5, where n can be up to 6. My first attempt was to write

f_with_norm_dependency = lambda unit_vec, norm: f(unit_vec * norm)
and then get the derivative by
f_with_norm_dependency_grad = jax.jacfwd(f_with_norm_dependency, argnums=1)

however, this is numerically unstable for higher-order derivatives. The reason seems to be that the norm function itself becomes unstable in this case, so

norm_with_norm_dependency = lambda norm: jnp.linalg.norm(unit_vec * norm)
norm_with_norm_dependency_grad5 = jax.jacfwd(jax.jacfwd(jax.jacfwd(jax.jacfwd(jax.jacfwd(norm_with_norm_dependency, argnums=0) ,argnums=0),argnums=0) ,argnums=0) ,argnums=0)

is unstable for small norms. Obviously, the first derivative is always 1, and all higher-order derivatives should always be 0. Therefore, I tried to write a custom gradient where I check if the tangent vector is colinear to the primal vector, and if it is the case, simply return the constant 1.0 like this:

@norm.defjvp
def norm_jvp(primals, tangents):
    rs, safe, axis = primals
    rs_dot, _, _ = tangents

    # Weird but numerically more stable
    quotient = rs / rs_dot
    is_colinear = (quotient - quotient.mean(axis=axis, keepdims=True)).std(axis=axis) < 1e-10
    
    # Less numerically stable
    # rs_normalized = rs / jnp.linalg.norm(rs, axis=axis, keepdims=True)
    # rs_dot_normalized = rs_dot / jnp.linalg.norm(rs_dot, axis=axis, keepdims=True)
    # is_colinear = jnp.allclose(rs_normalized, rs_dot_normalized, rtol=1e-5, atol=1e-8, equal_nan=True)

    tangent1 = jnp.sum(rs * rs_dot, axis=axis) / norm(rs, axis=axis)
    tangent2 = 1.0

    tangent = tangent2 * is_colinear + tangent1 * ~is_colinear

    return norm(rs), tangent

This works as long as I only use jax.jacfwd, but it breaks as soon as I try to combine it with jax.jacrev. I need to use both because the function takes a list of many different those r as input, and in another part of the code (a regularization loss), I need to calculate the Laplacian with respect to those r, which means I need to calculate first reverse than forward jacobians of f with respect to r itself, not |r|.
I tried to write a custom vjp for the jvp, but didn't get anything to work. Do you have an idea how this could be done?