NaN derivative in square(sqrt(x)) in model local variable

[bob-carpenter]

Simon Barthelmé reported via the stan-users list. I simplified the reported model to this minimal instance:

parameters {
  real x;
}
model {
  real y;
  y <- sqrt(square(x-x));

  x ~ normal(0,1);
}

As Simon correctly diagnosed, this model causes a problem with autodiff, as can be seen when running from RStan 1.3:

TESTING GRADIENT FOR MODEL 'barthelme2' NOW (CHAIN 4).

TEST GRADIENT MODE

 Log probability=-0.0571533

 param idx           value           model     finite diff           error
         0       -0.338093             nan        0.338093             nan

First step is to move this into a unit test that displays the behavior.

[bob-carpenter]

We can't really fix this problem without a huge amount of work moving to expression templates; with expression templates, we could follow Eigen in carrying out symbolic math at compile time, such as reducing sqrt(square(u)) to u.

I just checked that the boundary conditions in both sqrt() and square() are correct, as is the application of the chain rule.

The problem is that the chain rule is applied literally with floating-point arithmetic, which given

sqrt'(u) = 0.5 / sqrt(u)

square'(v) = 2 * v

yields

d/dx sqrt(square(x-x))

  = sqrt'(square(x-x)) * square'(x - x) * d/dx (x - x)

  = (+ infinity)  * 0 * 0

  = NaN.

I am going to add a discussion of this case to the manual, so I'm linking #

The problem is that (infinity * 0) produces NaN in floating-point arithmetic. There's no way given just the infinity and 0 that mulitplying them should produce 1.

[bob-carpenter]

Implement as Marcus suggested in e-mail:

Putting those two things together, maybe we can just rewrite chain() for sqrt() to be:

void chain() {
  if (val_ > 1e-300)
    avi_->adj_ += adj_ / (2.0 * val_);
  else
    avi_->adj_ += adj_ / 2.0;
  }

I asked: Will this do the right thing for HMC and BFGS? And Marcus responded:

This is effectively returning a subgradient for points where the gradient isn't defined. HMC won't care so long as same subgradient is always returned (which it is in this case). The regions where this will happen are extremely small, so, except for a few pathological cases, I wouldn't expect it to impact either method significantly.

[syclik]

This issue was moved to stan-dev/math#21