patrick-kidger / optimistix

Nonlinear optimisation (root-finding, least squares, ...) in JAX+Equinox. https://docs.kidger.site/optimistix/
Apache License 2.0
265 stars 12 forks source link

Error in "optimistix/docs/examples /optimise_diffeq.ipynb" #37

Open KhayrullevJokhongir opened 5 months ago

KhayrullevJokhongir commented 5 months ago

Hi, I was trying to run the example given in "optimistix/docs/examples /optimise_diffeq.ipynb". For some reason I am receiving error "--------------------------------------------------------------------------- ValueError Traceback (most recent call last) __ Cell 4 line 9 5 pred_values = batch_solve(parameters, y0s, saveat) 6 return values - pred_values ----> 9 (y0s_0, values_0) = get_data() 10 y0s = jnp.array(y0s_0) 11 values = jnp.array(values_0)

__ Cell 4 line 1 9 saveat = dfx.SaveAt(ts=jnp.linspace(0, 30, 20)) 10 batch_solve = eqx.filter_jit(eqx.filter_vmap(solve, in_axes=(None, 0, None))) ---> 11 values = batch_solve(true_parameters, y0s, saveat) 12 return y0s, values

[... skipping hidden 21 frame]

__ Cell 4 line 2 19 t1 = saveat.subs.ts[-1] 20 dt0 = 0.1 ---> 21 sol = dfx.diffeqsolve( 22 term, 23 solver, 24 t0, 25 t1, ... --> 305 raise ValueError("No arrays to thread error on to.") 306 dynamic_x = _error(dynamic_x, pred, index, msgs=msgs, on_error=on_error) 307 return combine(dynamic_x, static_x)

ValueError: No arrays to thread error on to.". I would appreciate if I could get help to fix it.

patrick-kidger commented 5 months ago

What version of Equinox and Diffrax do you have? This was a known crash when using a particular pair of versions, but should be fixed when using the latest version of both.