JAX：在回复完成之前取消了 execute_request 消息的未来

Question

我有一个优化问题，我正试图用牛顿法解决。为了计算雅可比矩阵，我使用 jax.jacobian.

我的目标函数叫做

calc_ms

。我用牛顿法求根：

def newton(f, x_0, tol=1e-5, max_iter=10):
    x = x_0
    q = lambda x: x - jnp.linalg.solve(jax.jacobian(f)(x), f(x))
    error = tol + 1
    n = 0
    while error > tol:
        n+=1
        if(n > max_iter):
            raise Exception('Max iteration reached without convergence')
        y = q(x)
        if(any(jnp.isnan(y))):
            raise Exception('Solution not found with NaN generated')
        error = jnp.linalg.norm(x - y)
        x = y
        print(f'iteration {n}, error = {error:.5f}')
    print('\n' + f'Result = {x} \n')
    return x

然后我用这个函数来解决问题：

newton(lambda delta: calc_ms(delta, dist_util, food_exp, market_shares, corr_mat, n_stores), delta0).block_until_ready()

一段时间后内核崩溃并出现此错误：

Canceled future for execute_request message before replies were done
The Kernel crashed while executing code in the the current cell or a previous cell. Please review the code in the cell(s) to identify a possible cause of the failure. Click here for more info. View Jupyter log for further details.