EnzymeAD / EnzymeAD/Enzyme

NaN Hessian for saturated `tanh` when computing Hessian

Open
#2,836 1 comment 1 reaction 2 assignees Claimed by @wsmoses View on GitHub
Dominant language
LLVM
Stars
1.7k
Forks
188
Avg merge
2d 4h
Merged PRs (30d)
22

Description

Enzyme returns `NaN` Hessian entries for a scalar function involving saturated `tanh`, even though the expected Hessian is finite and evaluates to zero in double precision.

This appears related to Enzyme's derivative rule for `tanh` (https://github.com/EnzymeAD/Enzyme/blob/main/enzyme/Enzyme/InstructionDerivatives.td#L354), which uses:

```
d/dx tanh(x) = 1 / (cosh(x) * cosh(x))
```
This is mathematically correct but numerically unstable, because cosh(x) grows very fast and for large |x|, computing the cosh^2(x) can overflow.

## Minimal Reproducer

```cpp
#include
#include
#include
#include
#include
#include

#include

extern int enzyme_dup;
extern int enzyme_dupnoneed;
extern int enzyme_out;
extern int enzyme_const;

template
RT __enzyme_autodiff(void*, Args...);

template
RT __enzyme_fwddiff(void*, Args...);

// Function under test.
void f(double* x, double* y) {
y[0] = std::tanh(x[0] / (2.0 * x[1])) + x[0];
}

template
void reverse_grad_func(double* x, double* dx, double* y, double* dy) {
__enzyme_autodiff(
(void*)F,
enzyme_dup, x, dx,
enzyme_dup, y, dy
);
}

template
void forward_over_reverse_func(
double* x,
double* dx_seed,
double* y,
double* dy,
double* grad_out,
double* hess_out
) {
__enzyme_fwddiff(
(void*)reverse_grad_func,
enzyme_dup, x, dx_seed,
enzyme_dup, grad_out, hess_out,
enzyme_const, y,
enzyme_const, dy
);
}

int main() {
std::vector x = {-0.62, -0.0008};

std::vector y(1, 0.0);
std::vector dy(1, 0.0);
std::vector dx_seed(2, 0.0);
std::vector grad_out(2, 0.0);
std::vector hess_out(2, 0.0);

std::vector> H(2, std::vector(2, 0.0));

for (size_t col = 0; col < 2; ++col) {
std::fill(y.begin(), y.end(), 0.0);
std::fill(dy.begin(), dy.end(), 0.0);
std::fill(dx_seed.begin(), dx_seed.end(), 0.0);
std::fill(grad_out.begin(), grad_out.end(), 0.0);
std::fill(hess_out.begin(), hess_out.end(), 0.0);

dy[0] = 1.0;
dx_seed[col] = 1.0;

forward_over_reverse_func(
x.data(),
dx_seed.data(),
y.data(),
dy.data(),
grad_out.data(),
hess_out.data()
);

for (size_t row = 0; row < 2; ++row) {
H[row][col] = hess_out[row];
}
}

std::cout << "H = [["
<< H[0][0] << ", " << H[0][1] << "], ["
<< H[1][0] << ", " << H[1][1] << "]]\n";
}
```

The output is: H = [[-nan, -nan], [-nan, -nan]].
Because the tanh argument is about 387.5, cosh(z) * cosh(z) overflows to inf, and caused an evaluation of inf/inf = NaN.

The expected output is H = [[0, 0], [0, 0]], which aligns with what I got from other AD tools, like PyTorch, JAX, CppAD, CasADi, and CoDiPack.

Has Enzyme considered using a more numerically stable derivative rule for tanh like:

`d/dx tanh(x) = 1 - tanh(x)^2`?

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.