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.