Replaced gpu<->cpu copies with host_callback in minimize.

scico/solver.py

@@ -63,6 +63,7 @@
 import numpy as np
 
 import jax
+import jax.experimental.host_callback as hcb
 
 import scico.numpy as snp
 from scico.numpy import BlockArray
@@ -210,11 +211,6 @@ def minimize(
     x0_shape = x0.shape
     x0_dtype = x0.dtype
     x0 = x0.ravel()  # if x0 is a BlockArray it will become a DeviceArray here
-    if isinstance(x0, jax.interpreters.xla.DeviceArray):
-        dev = x0.device_buffer.device()  # device for x0; used to put result back in place
-        x0 = np.array(x0).astype(float)
-    else:
-        dev = None
 
     # Run the SciPy minimizer
     if method in (
@@ -229,13 +225,28 @@ def minimize(
         min_func = _wrap_func(func_, x0_shape, x0_dtype)
         jac = False
 
-    res = spopt.minimize(
-        min_func,
-        x0=x0,
-        args=args,
-        jac=jac,
-        method=method,
-        options=options,
+    res = Optional[Any]
+
+    def fun(x0):
+        nonlocal res  # To use the external res and update side effect
+        res = spopt.minimize(
+            min_func,
+            x0=x0,
+            args=args,
+            jac=jac,
+            method=method,
+            options=options,
+        )  # Returns OptimizeResult
+        return res.x.astype(x0.dtype)  # Return for host_callback
+
+    if isinstance(x0, jax.interpreters.xla.DeviceArray):
+        dev = x0.device_buffer.device()  # Save to return x0 to its proper device later
+
+    # hcb call with side effects to get the OptimizeResult on the same device it was called
+    hcb.call(
+        fun,
+        arg=x0,
+        result_shape=x0,
     )
 
     # un-vectorize the output array, put on device
@@ -245,8 +256,8 @@ def minimize(
 
     res.x = res.x.astype(x0_dtype)
 
-    if dev:
-        res.x = jax.device_put(res.x, dev)
+    if isinstance(x0, jax.interpreters.xla.DeviceArray):
+        res.x = jax.device_put(res.x, dev)  # Return x0/res.x to its device
 
     if iscomplex:
         res.x = _join_real_imag(res.x)

scico/test/test_solver.py

@@ -179,6 +179,24 @@ def f(x):
     assert out.x.shape == x.shape
     np.testing.assert_allclose(out.x.ravel(), expected, rtol=5e-4)
 
+    # Check if minimize returns the object to the proper device
+    devices = jax.devices()
+
+    # For default device:
+    x0 = jax.device_put(snp.zeros_like(x), devices[0])
+    out = solver.minimize(f, x0=x0, method=method)
+    assert out.x.device_buffer.device() == devices[0]
+    assert out.x.shape == x0.shape
+    np.testing.assert_allclose(out.x.ravel(), expected, rtol=5e-4)
+
+    # If there are more than one device present:
+    if len(devices) > 1:
+        x0 = jax.device_put(snp.zeros_like(x), devices[1])
+        out = solver.minimize(f, x0=x0, method=method)
+        assert out.x.device_buffer.device() == devices[1]
+        assert out.x.shape == x0.shape
+        np.testing.assert_allclose(out.x.ravel(), expected, rtol=5e-4)
+
 
 def test_split_join_array():
     x, key = random.randn((4, 4), dtype=np.complex64)
@@ -206,22 +224,22 @@ def test_split_join_blockarray():
 
 
 def test_bisect():
-    f = lambda x: x**3
+    f = lambda x: x ** 3
     x, info = solver.bisect(f, -snp.ones((5, 1)), snp.ones((5, 1)), full_output=True)
     assert snp.sum(snp.abs(x)) == 0.0
     assert info["iter"] == 0
     x = solver.bisect(f, -2.0 * snp.ones((5, 3)), snp.ones((5, 3)), xtol=1e-5, ftol=1e-5)
     assert snp.max(snp.abs(x)) <= 1e-5
     assert snp.max(snp.abs(f(x))) <= 1e-5
     c, key = random.randn((5, 1), dtype=np.float32)
-    f = lambda x, c: x**3 - c**3
+    f = lambda x, c: x ** 3 - c ** 3
     x = solver.bisect(f, -snp.abs(c) - 1, snp.abs(c) + 1, args=(c,), xtol=1e-5, ftol=1e-5)
     assert snp.max(snp.abs(x - c)) <= 1e-5
     assert snp.max(snp.abs(f(x, c))) <= 1e-5
 
 
 def test_golden():
-    f = lambda x: x**2
+    f = lambda x: x ** 2
     x, info = solver.golden(f, -snp.ones((5, 1)), snp.ones((5, 1)), full_output=True)
     assert snp.max(snp.abs(x)) <= 1e-7
     x = solver.golden(f, -2.0 * snp.ones((5, 3)), snp.ones((5, 3)), xtol=1e-5)