NaN Hessian for saturated `tanh` when computing Hessian
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 <algorithm>
#include <chrono>
#include <cmath>
#include <cstddef>
#include <iostream>
#include <vector>
#include <enzyme/enzyme>
extern int enzyme_dup;
extern int enzyme_dupnoneed;
extern int enzyme_out;
extern int enzyme_const;
template <typename RT, typename... Args>
RT __enzyme_autodiff(void*, Args...);
template <typename RT, typename... Args>
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 <auto F>
void reverse_grad_func(double* x, double* dx, double* y, double* dy) {
__enzyme_autodiff<void>(
(void*)F,
enzyme_dup, x, dx,
enzyme_dup, y, dy
);
}
template <auto F>
void forward_over_reverse_func(
double* x,
double* dx_seed,
double* y,
double* dy,
double* grad_out,
double* hess_out
) {
__enzyme_fwddiff<void>(
(void*)reverse_grad_func<F>,
enzyme_dup, x, dx_seed,
enzyme_dup, grad_out, hess_out,
enzyme_const, y,
enzyme_const, dy
);
}
int main() {
std::vector<double> x = {-0.62, -0.0008};
std::vector<double> y(1, 0.0);
std::vector<double> dy(1, 0.0);
std::vector<double> dx_seed(2, 0.0);
std::vector<double> grad_out(2, 0.0);
std::vector<double> hess_out(2, 0.0);
std::vector<std::vector<double>> H(2, std::vector<double>(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<f>(
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`?
0 条评论