Repository: google/jax
Issue: Line-search x64 type promotion bug

There seems to be a bug in the line-search when enabling x64 mode but optimizing a purely float32 function.

```python
import jax.numpy as jnp
import jax.scipy.optimize

jax.config.update("jax_enable_x64", True)


def f(x):
    return jnp.sum(x ** 2)


x0 = jnp.zeros(2, dtype=jnp.float32)
jax.scipy.optimize.minimize(f, x0, method='BFGS')
```

```
TypeError: body_fun output and input must have identical types, got
_ZoomState(done=ShapedArray(bool[]), failed=ShapedArray(bool[]), j=ShapedArray(int64[], weak_type=True), a_lo=ShapedArray(float64[]), phi_lo=ShapedArray(float64[]), dphi_lo=ShapedArray(float64[]), a_hi=ShapedArray(float64[]), phi_hi=ShapedArray(float64[]), dphi_hi=ShapedArray(float64[]), a_rec=ShapedArray(float64[]), phi_rec=ShapedArray(float64[]), a_star=ShapedArray(float64[]), phi_star=ShapedArray(float64[]), dphi_star=ShapedArray(float64[]), g_star=ShapedArray(float64[2]), nfev=ShapedArray(int64[], weak_type=True), ngev=ShapedArray(int64[], weak_type=True))
and
_ZoomState(done=ShapedArray(bool[], weak_type=True), failed=ShapedArray(bool[], weak_type=True), j=ShapedArray(int64[], weak_type=True), a_lo=ShapedArray(float64[], weak_type=True), phi_lo=ShapedArray(float32[]), dphi_lo=ShapedArray(float64[]), a_hi=ShapedArray(float64[], weak_type=True), phi_hi=ShapedArray(float64[]), dphi_hi=ShapedArray(float64[]), a_rec=ShapedArray(float64[], weak_type=True), phi_rec=ShapedArray(float64[]), a_star=ShapedArray(float64[], weak_type=True), phi_star=ShapedArray(float32[]), dphi_star=ShapedArray(float64[]), g_star=ShapedArray(float32[2]), nfev=ShapedArray(int64[], weak_type=True), ngev=ShapedArray(int64[], weak_type=True)).
```

-> `g_star` type differs

Is this expected behavior or a bug?


Changed files:
tests/scipy_optimize_test.py, jax/_src/scipy/optimize/bfgs.py

diff --git a/jax/_src/scipy/optimize/bfgs.py b/jax/_src/scipy/optimize/bfgs.py
index 7d7a8aff4..197e3844c 100644
--- a/jax/_src/scipy/optimize/bfgs.py
+++ b/jax/_src/scipy/optimize/bfgs.py
@@ -144,7 +144,7 @@ def minimize_bfgs(
     rho_k = jnp.reciprocal(_dot(y_k, s_k))
 
     sy_k = s_k[:, jnp.newaxis] * y_k[jnp.newaxis, :]
-    w = jnp.eye(d) - rho_k * sy_k
+    w = jnp.eye(d, dtype=rho_k.dtype) - rho_k * sy_k
     H_kp1 = (_einsum('ij,jk,lk', w, state.H_k, w)
              + rho_k * s_k[:, jnp.newaxis] * s_k[jnp.newaxis, :])
     H_kp1 = jnp.where(jnp.isfinite(rho_k), H_kp1, state.H_k)
diff --git a/tests/scipy_optimize_test.py b/tests/scipy_optimize_test.py
index 7d6f02177..562168a0e 100644
--- a/tests/scipy_optimize_test.py
+++ b/tests/scipy_optimize_test.py
@@ -71,8 +71,8 @@ class TestBFGS(jtu.JaxTestCase):
     {"testcase_name": f"_func={func_and_init[0].__name__}_maxiter={maxiter}",
      "maxiter": maxiter, "func_and_init": func_and_init}
     for maxiter in [None]
-    for func_and_init in [(rosenbrock, np.zeros(2)),
-                          (himmelblau, np.ones(2)),
+    for func_and_init in [(rosenbrock, np.zeros(2, dtype='float32')),
+                          (himmelblau, np.ones(2, dtype='float32')),
                           (matyas, np.ones(2) * 6.),
                           (eggholder, np.ones(2) * 100.)]))
   def test_minimize(self, maxiter, func_and_init):