I have an optimization problem which I am trying to solve with Newton's method. To calculate jacobian matrix, I use jax.jacobian.
My objective function is called calc_ms. I use Newton's method to find roots:
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
And then I use this function to solve the problem:
newton(lambda delta: calc_ms(delta, dist_util, food_exp, market_shares, corr_mat, n_stores), delta0).block_until_ready()
After some time kernel crashed with this error:
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.
