Overloading `jax` operations to avoid undefined derivatives

[alleSini99]

Hello,
I would like to overload some jax operations to work with cases where the derivative of a function f is undefined/diverging, but the derivative of a composite function of f, let's say g(f), is defined.
I report here an example. In this case f is -inf and so exp(f) is 0. If I compute the derivative of f with jax.grad this is inf as expected, so the derivative of exp(f) is nan since this is computed from the chain rule as d exp(f) = exp(f) df and I get a 0 times inf. However, thanks to mathematical simplifications, the correct d exp(f) is finite and I would like to obtain this as the result. We have partially solved the problem by defining special versions of log and exp based on equinox classes, but is there a general way to handle this issue for arbitrary jax operations within quad?
Example:

import jax.numpy as jnp 
import jax 
import equinox as eqx

def psii(theta, x): 
    return theta * x
thetas = jnp.array([1., 1.])
def f(theta, x): 
    return jnp.log(psii(theta[0], x) - psii(theta[1], x))
expf = lambda thetas, x: jnp.exp(f(thetas, x))
x = 0.5
print("f:", f(thetas, x))
print("expf:", expf(thetas, x))
print("df:", jax.grad(f)(thetas, x))
print("d expf:", jax.grad(expf)(thetas, x))

class MagicLog(eqx.Module):
    val : jax.Array
    def __init__(self, x):
        self.val = x
    def __jax_array__(self, dtype=None): 
        return jnp.log(self.val)
def magiclog(x):
    return MagicLog(x)
def magicexp(x):
    if isinstance(x, MagicLog):
        return x.val
    else:
        return jnp.exp(x)
def magiclogpsi(theta, x): 
    return magiclog(psii(theta[0], x) - psii(theta[1], x))
magicexpf = lambda thetas, x: magicexp(magiclogpsi(thetas, x))
print("magic expf:", magicexpf(thetas, x))
print("magic d expf:", jax.grad(magicexpf)(thetas, x))

@PhilipVinc

[patrick-kidger]

I think you probably want to wrap your operations in a jax.custom_jvp. This kind of thing isn't really related to either Equinox or a Quax. :)

Side note, I can see you're using __jax_array__. I'd recommend not using this -- it was an early experiment, deliberately undocumented, and is not well-supported throughout JAX.

[PhilipVinc]

@patrick-kidger I'm not sure that's exactly what we want to do...

We want to have some exp and log() functions that analytically cancel each other (in general, but especially when doing AD).

Right now, if we do AD through $$g(x)=\exp(\log(f(x)) $$ we get something that is $$\nabla g(x) = \exp(\log(f(x)) \cdot \nabla \log(f(x)) = \exp(\log(f(x)) \cdot \frac{\nabla f(x)}{f(x)} $$.

This is analytically equivalent to $$\nabla g(x)=\nabla f(x)$$ if $$f(x) \neq0$$.

But if $$f(x)=0$$, Jax does the chain rule so multiplies and divide by $$f(x)$$ and so we get a nan, because we divide by zero.
My idea was that by creating a custom 'symbolic-log' type we could symbolically eliminate exp(log(f(x)) before we lower to jaxpr, pretty much like the analytical zero example of quax propagates symbolically zeros before they are lowered in the jaxpr.

We'd be happy to use jax.custom_jvp/vjp but I'm unsure how, in this case? even if we define a special 'exp' function with a custom AD rule, we cannot check whether the internal function is the logarithm and do the simplification analytically?

[patrick-kidger]

Ah, okay! That makes sense.

So I think hoping to do things before you hit the jaxpr is likely to be tricky. That's just not going to compose well in all kinds of ways.

My first thought for how you might try and tackle this is to write a custom jaxpr-to-jaxpr function (directly, not using Quax), that checks for the appropriate pattern and cancels things as appropriate. In compiler language this is basically a peephole optimization (aka algebraic pattern matching).

That would need a bit of care to still do the cancellation when you have, for example, a log at the start of a loop body and an exp at an end, but I think should still be essentially doable.

I hope that helps!

[nstarman]

is there a general way to handle this issue for arbitrary jax operations within quax?

Yes, and this is a good use of Quax. There is now a how-to page for it: Recover a gradient JAX computes as NaN.

The idea is the one behind your MagicLog, but expressed as a Value with dispatch rules instead of __jax_array__: give the singular operation a type that stands for its result without computing it, and register the cancelling operation to unwrap it. The singular step then never appears in the program being differentiated, so no inf is ever formed to multiply by zero.

Reduced to the smallest case:

def g(a, b):
    return jnp.exp(jnp.log(a - b))

jax.grad(g)(1.0, 1.0)  # nan

class Log(quax.ArrayValue):
    """Stands for log(arg), left unevaluated."""
    arg: jax.Array = eqx.field(converter=jnp.asarray)
    def aval(self): return jax.typeof(self.arg)
    def materialise(self): return jnp.log(self.arg)

@quax.register(jax.lax.log_p)
def log_tracked(a: Tracked, **kw) -> Log:
    return Log(a.array)

@quax.register(jax.lax.exp_p)
def exp_log(a: Log, **kw) -> Tracked:
    return Tracked(a.arg)          # exp(log(y)) is y, exactly

jax.grad(quax_g)(1.0, 1.0)  # 1.0

Nothing is clipped or approximated — exp(log(y)) is y — and results away from the singularity are unchanged (g(3, 1) == 2.0, gradient 1.0). The page has the full runnable version; it is executed as part of the test suite, so it stays correct.

Two caveats that cost me time when reproducing your example, both on the page:

1. Every primitive between your inputs and the cancellation needs a rule. Miss one and it falls back to materialise, your type quietly becomes a plain array, and log is evaluated the ordinary way — you get the same nan back with nothing to indicate the rules never fired. In your example theta[0] is the culprit: it lowers to slice_p followed by squeeze_p, not one primitive. Raising from materialise instead of returning the array turns that silent failure into a loud one.

2. The pattern needs an exact inverse pair that cancels symbolically — exp/log, sqrt/square, or one of your own. Where the limit is finite but the expression does not cancel, jax.custom_jvp is the right tool instead.

Your original case gives [0.5, -0.5] under this treatment, against [nan, nan] from plain JAX.

---

Written by Claude, reviewed and sent by @nstarman.