AD in ill-conditioned z/√(s²+z²) without custom rules

Specific Domains Numerics
, ,
1

I would like to implement a z/sqrt(s^2+z^2) that

  1. works for when z and s have wildly different magnitudes, but are both finite,
  2. also gives the correct result when AD’s via the most common frameworks.

I would prefer to avoid writing custom rules if possible. This part makes the exercise tricky.

Here is an MWE of failures using Enzyme as a demo because it traps NaNs. I would appreciate suggestions on how to do this. Or maybe someone convincing me that this is not possible without custom rules.

using Enzyme, Test

Enzyme.Compiler.CheckNan[] = true # error for NaNs

function _z_over_hypot(z, s)
    r = z / s
    if abs(r) ≤ 1
        r / sqrt(one(r) + r^2)
    else
        u = inv(r)
        sign(r) / sqrt(one(u) + u^2)
    end
end

@testset "z over hypot torture test" begin
    # V = [1.0, 2.0, 3.0] # for simple tests
    V = [1e-180, -1e-180, 0.5e-100, 1e-100, 2e-100, 1e-50, 1e100, 1e200]
    for s in V
        for z in V
            for dz in V
                @show s, z, dz
                dx, x = Enzyme.autodiff(Enzyme.ForwardWithPrimal, z -> SK._z_over_hypot(z, s),
                                        Enzyme.Duplicated(z, dz))
                _z = BigFloat(z)
                _s = BigFloat(s)
                _r = _z / _s
                _q = 1 + _r^2
                _x = z / √(z^2 + s^2)
                _dx = 1 / _s / (_q * √_q)
                @test x ≈ Float64(_x)
                @test dx ≈ Float64(_dx) * dz
            end
        end
    end
end
2

What’s wrong with the naive _z_over_hypot(z, s) = z / hypot(z,s)?

I haven’t tried Enzyme, but a simplified version of your test that simply checks whether it matches the BigFloat result passes:

julia> @testset "z over hypot torture test" begin
           # V = [1.0, 2.0, 3.0] # for simple tests
           V = [1e-180, -1e-180, 0.5e-100, 1e-100, 2e-100, 1e-50, 1e100, 1e200]
           for s in V
               for z in V
                   @test _z_over_hypot(z,s) ≈ Float64(_z_over_hypot(big(z),big(s)))
               end
           end
       end
Test Summary:             | Pass  Total  Time
z over hypot torture test |   64     64  0.0s
3

It fails for AD. That’s why Enzyme is in the tests MWE.

4

I see. I guess the basic problem here is that the naive quotient rule for \partial/\partial z of this function (which is what AD will do by default, probably) gives:

\frac{\partial}{\partial z} \frac{z}{\operatorname{hypot}(z,s)} = \frac{\operatorname{hypot}(z,s) - \frac{z^2}{\operatorname{hypot}(z,s)}}{\operatorname{hypot}(z,s)^2} = \cdots = \frac{s^2}{\operatorname{hypot}(z,s)^{3}}

which suffers catastrophic cancellation in the naive numerator for |z| \gg |s| unless you analytically simplify it, which AD won’t do (unless you re-arrange the input expression). And even the simplified expression with s^2 in the numerator, which has no catastrophic cancellation, is susceptible to spurious over/underflow if both s and z are very large or small (that is, the numerator and denominator can separately under/overflow, even if the final answer would be in the representable range).

5

I suspect that the overflow/underflow issue, in particular, will be hard to solve without a custom rule — it’s really hard to avoid spurious underflow of intermediate terms in expressions like this unless you have control over the arrangements of the terms. Your control over the intermediate terms that AD generates seems like it will inevitably be pretty fragile — even if you find a formulation that works, an update to the AD system could easily break it.

If you don’t have to worry about over/underflow, e.g. if your magnitudes are at worst 10^{\pm 50} or so, then the problem becomes much easier.

6

I’ve run into this kind of thing a lot when writing Quaternionic.jl and SphericalFunctions.jl, and have come around to the idea that AD is perfect for quick-and-dirty calculations, but it inherently changes the numerics because it actually has to compute different things. So the optimal approach would frequently use a different one from what AD uses, which means you have to be careful about your results. (Chris Rackauckas has made this point a number of times, more extensively.)

But in this case, I think getting around the over/underflow issues specifically isn’t too hard in a way that’s AD-compatible: use the rescale-by-powers-of-two trick:

function z_over_hypot(z, s)
    e = exponent(max(abs(z), abs(s)))
    z, s = ldexp(z, -e), ldexp(abs(s), -e)   # now max(|z|, s) ∈ [1, 2)
    if abs(z) ≤ s
        r = z / s
        r / sqrt(one(r) + r^2)
    else
        u = s / z
        sign(z) / sqrt(one(u) + u^2)
    end
end

The only cases of yours where this gives “bad” results are where the exact derivative itself overflows (|z|=|s|=1e-180, dz=1e200; derivatives should be ±Inf but result is NaN) or is subnormal (various combinations; derivatives have relative errors like 1e-7).

7

Oh, also, I’ve given in a lot lately and just “written” rules. Luckily, Claude is great at writing all the magical invocations needed for all the different packages, so I’ve just given over the AD-rules responsibility to Claude.

8

As @stevengj suspected, that in itself is not enough to fix the AD issue for me. But the ldexp approach works fine with a custom AD rule.

Yes, LLMs are nice for writing AD rule boilerplate once I implement the partials, but for simple scalar cases I also just copy and modify existing code.

9

I think it does actually fix it, even without a custom rule — except for the few noted cases. But there are some problems in your test code:

  1. your _dx formula is wrong for negative _s,
  2. it doesn’t compute the reference _x in BigFloat, and
  3. it stops on the first error, which makes it look like everything is wrong.

If I just run your original testset with my version of z_over_hypot, I get

7 passed, 7 failed, 1 errored, 0 broken.

But if I fix your testset as below, and run with my version of z_over_hypot, I get

1016 passed, 0 failed, 4 errored, 0 broken.

Those four errors are the ones with |s| = 1.0e-180, |z| = 1.0e-180, dz = 1.0e200, where we get NaN instead of the correctly overflowed \pm Inf.

In any case, that’s a big improvement over your original _z_over_hypot, which with the corrected testset reports

733 passed, 115 failed, 88 errored, 0 broken.

Here’s the revised test code:

using Enzyme, Test

Enzyme.Compiler.CheckNan[] = true # error for NaNs

@testset "z over hypot torture test" begin
    V = [1e-180, -1e-180, 0.5e-100, 1e-100, 2e-100, 1e-50, 1e100, 1e200]
    # Use nested testsets so we keep running after errors
    @testset "s = $s, z = $z, dz = $dz" for s in V, z in V, dz in V
        dx, x = Enzyme.autodiff(Enzyme.ForwardWithPrimal, z -> z_over_hypot(z, s),
                                Enzyme.Duplicated(z, dz))
        _z = BigFloat(z)
        _s = BigFloat(s)
        _h2 = _z^2 + _s^2
        _x = _z / √_h2  # This is now BigFloat
        _dx = _s^2 / (_h2 * √_h2) * dz  # This formula works with negative _s
        @test x ≈ Float64(_x)
        @test dx ≈ Float64(_dx)
    end
end