Repository: google/jax

Issue:
jax.scipy.optimize hangs
I'm comparing performance for the scipy BFGS implementation and the replica in jax, but the one in jax keeps hanging and/or being 100x slower. To make sure this is not only my system, I have verified that the code below hangs on Colab.
```python
import jax
import jax.numpy as jnp
import jax.scipy.optimize
import jax.random

jX = jnp.arange(-2, 8, 0.03) # A 1D evaluation grid
targetp = jnp.array(
    [1.3615158,  3.1700504,  0.3901142,  0.8582838,
     -0.24486497,  3.078153,  1.2610698,  2.276195,
     0.84759176,  1.563921, -0.91339713,  0.23969328]) # The "true" parameters

def fun(p): # This function was once auto-generated, so it looks silly
    arg1 = p[2] + p[3] * jnp.tanh(p[4]  + p[5] * jX)
    arg2 = p[6] + p[7] * jnp.log(jnp.abs(p[8] + p[9] * jnp.arctan(p[10] + p[11] * jX)))
    return p[0] + p[1] * arg1 * arg2

targety = fun(targetp)
def mse_loss(p):
    return ((fun(p) - targety) ** 2).mean()
```

I can optimize `mse_loss` just fine in scipy.optimize.minimize, but when I run the code below, the interpreter hangs on the second iteration:
```python
key = jax.random.PRNGKey(1337)
keys = jax.random.split(key, 10)

for k in keys: # Try optimize from different starting points
    x0 = targetp + jax.random.normal(k)
    optx = jax.scipy.optimize.minimize(mse_loss, x0, method='BFGS')
    print(optx.fun, optx.nfev)
```

I have verified this with and without GPU and with and without jitted functions.

[This](https://colab.research.google.com/drive/1sBwLh5w_XRDABHJB2SrmwZtOTXu33EOD) is a link to the above code in Colab.
jax.scipy.optimize.minimize convergence problem
I attempted to minimize a simple test function with scipy.optimize.minimize and with jax.scipy.optimize.minimize. I used identical parameters and start point. scipy.optimize.minimize converged to the function minimum but jax.scipy.optimize.minimize did not.

```python
import time
import scipy.optimize
import jax
import jax.numpy as jnp
import jax.scipy.optimize
import jax.config
import autograd.numpy as anp
import autograd

jax.config.update("jax_enable_x64", True)

def zakharovFromIndices(x, ii):
    sum1 = (x**2).sum()
    sum2 = (0.5*ii*x).sum()
    answer = sum1+sum2**2+sum2**4
    return answer

def zakharov_jaxNumpy(x):
    ii = jnp.arange(1, len(x)+1, step=1)
    answer = zakharovFromIndices(x=x, ii=ii)
    return answer

def zakharov_autogradNumpy(x):
    ii = anp.arange(1, len(x)+1, step=1)
    answer = zakharovFromIndices(x=x, ii=ii)
    return answer

jEvalFunc = jax.jit(zakharov_jaxNumpy)
aEvalFunc = zakharov_autogradNumpy
aGradFunc = autograd.grad(aEvalFunc)

x0 = [600.0, 700.0, 200.0, 100.0, 90.0, 1e4]
toleranceChange = 1e-9
maxIter = 10000
jx0 = jnp.array(x0)
ax0 = anp.array(x0)

aOptimRes_x0 = scipy.optimize.minimize(fun=aEvalFunc, x0=ax0, method='BFGS', jac=aGradFunc)
jOptimRes_x0 = jax.scipy.optimize.minimize(fun=jEvalFunc, x0=jx0, method='BFGS')

print("scipy.optimize converged?: {}".format(aOptimRes_x0.fun<1e-6))
print("jax.scipy.optimize converged?: {}".format(jOptimRes_x0.fun<1e-6))
```


Discussion:

I tried to scale the Hessian matrix at the first iteration with the method proposed in Nocedal and Wright in "Numerical Optimization p.143 formula (6.20)" but it doesn't solve the issue. Similar problems have been mentioned in #5139. The Scipy implementation uses a second line search when the first one don't succeed which could explain the performance differences. 

Resolving patch:
diff --git a/jax/_src/scipy/optimize/bfgs.py b/jax/_src/scipy/optimize/bfgs.py
--- a/jax/_src/scipy/optimize/bfgs.py
+++ b/jax/_src/scipy/optimize/bfgs.py
@@ -55,6 +55,7 @@ class _BFGSResults(NamedTuple):
   f_k: jnp.ndarray
   g_k: jnp.ndarray
   H_k: jnp.ndarray
+  old_old_fval: jnp.ndarray
   status: Union[int, jnp.ndarray]
   line_search_status: Union[int, jnp.ndarray]
 
@@ -108,6 +109,7 @@ def minimize_bfgs(
       f_k=f_0,
       g_k=g_0,
       H_k=initial_H,
+      old_old_fval=f_0 + jnp.linalg.norm(g_0) / 2,
       status=0,
       line_search_status=0,
   )
@@ -124,6 +126,7 @@ def body_fun(state):
         state.x_k,
         p_k,
         old_fval=state.f_k,
+        old_old_fval=state.old_old_fval,
         gfk=state.g_k,
         maxiter=line_search_maxiter,
     )
@@ -153,7 +156,8 @@ def body_fun(state):
         x_k=x_kp1,
         f_k=f_kp1,
         g_k=g_kp1,
-        H_k=H_kp1
+        H_k=H_kp1,
+        old_old_fval=state.f_k,
     )
     return state
 
diff --git a/jax/_src/scipy/optimize/line_search.py b/jax/_src/scipy/optimize/line_search.py
--- a/jax/_src/scipy/optimize/line_search.py
+++ b/jax/_src/scipy/optimize/line_search.py
@@ -104,15 +104,16 @@ def _zoom(restricted_func_and_grad, wolfe_one, wolfe_two, a_lo, phi_lo,
   def body(state):
     # Body of zoom algorithm. We use boolean arithmetic to avoid using jax.cond
     # so that it works on GPU/TPU.
+    dalpha = (state.a_hi - state.a_lo)
     a = jnp.minimum(state.a_hi, state.a_lo)
     b = jnp.maximum(state.a_hi, state.a_lo)
-    dalpha = (b - a)
     cchk = delta1 * dalpha
     qchk = delta2 * dalpha
 
     # This will cause the line search to stop, and since the Wolfe conditions
     # are not satisfied the minimization should stop too.
-    state = state._replace(failed=state.failed | (dalpha <= 1e-10))
+    threshold = jnp.where((jnp.finfo(dalpha).bits < 64), 1e-5, 1e-10)
+    state = state._replace(failed=state.failed | (dalpha <= threshold))
 
     # Cubmin is sometimes nan, though in this case the bounds check will fail.
     a_j_cubic = _cubicmin(state.a_lo, state.phi_lo, state.dphi_lo, state.a_hi,
@@ -169,9 +170,9 @@ def body(state):
             hi_to_lo,
             state._asdict(),
             dict(
-                a_hi=a_lo,
-                phi_hi=phi_lo,
-                dphi_hi=dphi_lo,
+                a_hi=state.a_lo,
+                phi_hi=state.phi_lo,
+                dphi_hi=state.dphi_lo,
                 a_rec=state.a_hi,
                 phi_rec=state.phi_hi,
             ),
@@ -191,6 +192,9 @@ def body(state):
         ),
     )
     state = state._replace(j=state.j + 1)
+    # Choose higher cutoff for maxiter than Scipy as Jax takes longer to find
+    # the same value - possibly floating point issues?
+    state = state._replace(failed= state.failed | state.j >= 30)
     return state
 
   state = while_loop(lambda state: (~state.done) & (~pass_through) & (~state.failed),
@@ -213,7 +217,6 @@ class _LineSearchState(NamedTuple):
   phi_star: Union[float, jnp.ndarray]
   dphi_star: Union[float, jnp.ndarray]
   g_star: jnp.ndarray
-  saddle_point: Union[bool, jnp.ndarray]
 
 
 class _LineSearchResults(NamedTuple):
@@ -269,6 +272,11 @@ def restricted_func_and_grad(t):
   else:
     phi_0 = old_fval
     dphi_0 = jnp.dot(gfk, pk)
+  if old_old_fval is not None:
+    candidate_start_value = 1.01 * 2 * (phi_0 - old_old_fval) / dphi_0
+    start_value = jnp.where(candidate_start_value > 1, 1.0, candidate_start_value)
+  else:
+    start_value = 1
 
   def wolfe_one(a_i, phi_i):
     # actually negation of W1
@@ -292,18 +300,12 @@ def wolfe_two(dphi_i):
       phi_star=phi_0,
       dphi_star=dphi_0,
       g_star=gfk,
-      saddle_point=False,
   )
 
   def body(state):
     # no amax in this version, we just double as in scipy.
     # unlike original algorithm we do our next choice at the start of this loop
-    a_i = jnp.where(state.i == 1, 1., state.a_i1 * 2.)
-    # if a_i <= 0 then something went wrong. In practice any really small step
-    # length is a failure. Likely means the search pk is not good, perhaps we
-    # are at a saddle point.
-    saddle_point = a_i < 1e-5
-    state = state._replace(failed=saddle_point, saddle_point=saddle_point)
+    a_i = jnp.where(state.i == 1, start_value, state.a_i1 * 2.)
 
     phi_i, dphi_i, g_i = restricted_func_and_grad(a_i)
     state = state._replace(nfev=state.nfev + 1,
@@ -384,25 +386,28 @@ def body(state):
                      state)
 
   status = jnp.where(
-      state.failed & (~state.saddle_point),
+      state.failed,
       jnp.array(1),  # zoom failed
-      jnp.where(
-          state.failed & state.saddle_point,
-          jnp.array(2),  # saddle point reached,
           jnp.where(
               state.i > maxiter,
               jnp.array(3),  # maxiter reached
               jnp.array(0),  # passed (should be)
           ),
-      ),
   )
+  # Step sizes which are too small causes the optimizer to get stuck with a
+  # direction of zero in <64 bit mode - avoid with a floor on minimum step size.
+  alpha_k = state.a_star
+  alpha_k = jnp.where((jnp.finfo(alpha_k).bits != 64)
+                    & (jnp.abs(alpha_k) < 1e-8),
+                      jnp.sign(alpha_k) * 1e-8,
+                      alpha_k)
   results = _LineSearchResults(
       failed=state.failed | (~state.done),
       nit=state.i - 1,  # because iterations started at 1
       nfev=state.nfev,
       ngev=state.ngev,
       k=state.i,
-      a_k=state.a_star,
+      a_k=alpha_k,
       f_k=state.phi_star,
       g_k=state.g_star,
       status=status,


Test patch:
diff --git a/tests/scipy_optimize_test.py b/tests/scipy_optimize_test.py
--- a/tests/scipy_optimize_test.py
+++ b/tests/scipy_optimize_test.py
@@ -57,6 +57,13 @@ def func(p):
   return func
 
 
+def zakharovFromIndices(x, ii):
+  sum1 = (x**2).sum()
+  sum2 = (0.5*ii*x).sum()
+  answer = sum1+sum2**2+sum2**4
+  return answer
+
+
 class TestBFGS(jtu.JaxTestCase):
 
   @parameterized.named_parameters(jtu.cases_from_list(
@@ -94,6 +101,37 @@ def f(x):
     results = jax.scipy.optimize.minimize(f, jnp.ones(n), method='BFGS')
     self.assertAllClose(results.x, jnp.zeros(n), atol=1e-6, rtol=1e-6)
 
+  @jtu.skip_on_flag('jax_enable_x64', False)
+  def test_zakharov(self):
+    def zakharov_fn(x):
+      ii = jnp.arange(1, len(x) + 1, step=1)
+      answer = zakharovFromIndices(x=x, ii=ii)
+      return answer
+
+    x0 = jnp.array([600.0, 700.0, 200.0, 100.0, 90.0, 1e4])
+    eval_func = jax.jit(zakharov_fn)
+    jax_res = jax.scipy.optimize.minimize(fun=eval_func, x0=x0, method='BFGS')
+    self.assertLess(jax_res.fun, 1e-6)
+
+  def test_minimize_bad_initial_values(self):
+    # This test runs deliberately "bad" initial values to test that handling
+    # of failed line search, etc. is the same across implementations
+    initial_value = jnp.array([92, 0.001])
+    opt_fn = himmelblau(jnp)
+    jax_res = jax.scipy.optimize.minimize(
+        fun=opt_fn,
+        x0=initial_value,
+        method='BFGS',
+    ).x
+    scipy_res = scipy.optimize.minimize(
+        fun=opt_fn,
+        jac=jax.grad(opt_fn),
+        method='BFGS',
+        x0=initial_value
+    ).x
+    self.assertAllClose(scipy_res, jax_res, atol=2e-5, check_dtypes=False)
+
+
   def test_args_must_be_tuple(self):
     A = jnp.eye(2) * 1e4
     def f(x):