Computationally stable log-softplus

[KFrank]

Hi Steve!

I haven’t worked through this in detail, but a first look suggests that
the softplus() fix in the recent nightlies addresses your issue.

@stevethesteve

yields

1.8.1
diffA.abs().max() = tensor(8.4096e-07, dtype=torch.float64, grad_fn=<MaxBackward1>)
diffB.abs().max() = tensor(8.4096e-07, dtype=torch.float64, grad_fn=<MaxBackward1>)
diffB.abs().argmax() = tensor(61)
x[diffB.abs().argmax()] = tensor(-13.9000, dtype=torch.float64, grad_fn=<SelectBackward>)
diffgA.abs().max() = tensor(1.0000, dtype=torch.float64)
diffgB.abs().max() = tensor(0.1339, dtype=torch.float64)
diffgB.abs().argmax() = tensor(53)
x[diffgB.abs().argmax()] = tensor(-14.7000, dtype=torch.float64, grad_fn=<SelectBackward>)
lspgB[diffgB.abs().argmax()], lspg[diffgB.abs().argmax()] = tensor(0.8661) tensor(1.0000, dtype=torch.float64)

Running the script you posted on today’s nightly, 1.9.0.dev20210504,
seems to show that the issue is gone:

1.9.0.dev20210504
diffA.abs().max() = tensor(8.4096e-07, dtype=torch.float64, grad_fn=<MaxBackward1>)
diffB.abs().max() = tensor(8.4096e-07, dtype=torch.float64, grad_fn=<MaxBackward1>)
diffB.abs().argmax() = tensor(61)
x[diffB.abs().argmax()] = tensor(-13.9000, dtype=torch.float64, grad_fn=<SelectBackward>)
diffgA.abs().max() = tensor(1.6522e-07, dtype=torch.float64)
diffgB.abs().max() = tensor(1.6522e-07, dtype=torch.float64)
diffgB.abs().argmax() = tensor(162)
x[diffgB.abs().argmax()] = tensor(-3.8000, dtype=torch.float64, grad_fn=<SelectBackward>)
lspgB[diffgB.abs().argmax()], lspg[diffgB.abs().argmax()] = tensor(0.9890) tensor(0.9890, dtype=torch.float64)

Best.

K. Frank