tfp.optim.lbfgs_minimize fails where scipy.optimize and torch.optim.LBFGS succeed for poisson optimization

[lukasheinrich]

For this minimal example of fitting a 1-parameter poissonian likelihood Pois(10|par) where you would expect the best-fit value to be 10.0 tfp seems to fail:

tf.reset_default_graph()
sess = tf.Session()
def func_tf(pars):
    nll = -tfp.distributions.Poisson(pars).log_prob(10.)
    grad = tf.identity(tf.gradients(nll,pars))[0]
    return nll[0],grad

result = sess.run(tfp.optimizer.lbfgs_minimize(func_tf,tf.convert_to_tensor([1.6])))
'converged: {}, failed: {}, position: {}'.format(result.converged,result.failed,result.position)

gives:

'converged: False, failed: True, position: [11.124551]'

the optimization fails.

Interestingly it seems to be quite sensitive to the initial position. If I change the initial position to 10.5 the minimization succeeds:

result = sess.run(tfp.optimizer.lbfgs_minimize(func_tf,tf.convert_to_tensor([10.5])))
'converged: {}, failed: {}, position: {}'.format(result.converged,result.failed,result.position)

'converged: True, failed: False, position: [10.]'

As a cross-check I tried using the LBFGS optimizers in scipy.optimize and PyTorch to see if this is an issue, and in those cases the initial position 1.6 seems ok

Note: in scipy I had to cast to double precision due to: scipy/scipy#5832

I'm attaching a screentshot of the full notebook, which is available here: https://gist.github.com/lukasheinrich/327ab8bf1b96e9af3e483a4b26d435ea

cc @matthewfeickert @kratsg

[junpenglao]

This is a weird one - seems there is a region that it will fail to converge:

import functools

def _make_val_and_grad_fn(value_fn):
  @functools.wraps(value_fn)
  def val_and_grad(x):
    return tfp.math.value_and_gradient(value_fn, x)
  return val_and_grad

@_make_val_and_grad_fn
def func_tf(pars):
  return -tf.squeeze(tfd.Poisson(pars).log_prob(10.))

init_conditions = np.linspace(.1, 20., 100)
result = tfp.optimizer.lbfgs_minimize(
    func_tf, 
    init_conditions[..., tf.newaxis],
    stopping_condition=tfp.optimizer.converged_all,
    )

failed_not_converge = ~result.converged & result.failed
plt.plot(init_conditions, result.position)
plt.plot(init_conditions[failed_not_converge], 
         result.position[failed_not_converge], 'o', color='r');