Test for svmbir interface hang on GPU device

[bwohlberg]

The test functions test_prox and test_prox_weights in scico/test/linop/test_radon_svmbir.py hang when run on a GPU. A temporary workaround using @pytest.mark.skipif will be implemented, but the cause of the problem needs to be determined and addressed.

[bwohlberg]

The tests in question call scico.test.functional.prox_test, which calls scico.test.functional.prox_solve,
which calls scico.solver.minimize

scico/scico/test/functional/prox.py

Lines 14 to 25 in 090e607

	def prox_solve(v, v0, f, alpha):
	"""Evaluate the alpha-scaled proximal operator of f at v, using v0 as an
	initial point for the optimization."""
	fnc = lambda x: prox_func(x, v, f, alpha)
	fmn = minimize(
	fnc,
	v0,
	method="Nelder-Mead",
	options={"maxiter": 1000, "xatol": 1e-9, "fatol": 1e-9},
	)
	return fmn.x.reshape(v.shape), fmn.fun

When run on a GPU, the code appears to hang at the jax.experimental.host_callback call in scico.solver.minimize

scico/scico/solver.py

Lines 243 to 248 in 090e607

	# HCB call with side effects to get the OptimizeResult on the same device it was called
	res.x = hcb.call(
	fun,
	arg=x0,
	result_shape=x0, # From Jax-docs: This can be an object that has .shape and .dtype attributes
	)

[bwohlberg]

Closing since the problem seems to be within the jax host callback implementation.