Repository navigation
esn_forecast not jittable #60
Description
When trying to compile the forecast function, I run into the following error:
import orc
import jax
# integrate the Lorenz system
U,t = orc.data.lorenz63(tN=100, dt=0.01)
# train-test split
test_perc = 0.2
split_idx = int((1 - test_perc) * U.shape[0])
U_train = U[:split_idx, :]
t_train = t[:split_idx]
U_test = U[split_idx:, :]
t_test = t[split_idx:]
# Initialize and train the ESN
esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=400)
esn, R = orc.forecaster.train_RCForecaster(esn, U_train)
# Forecast!
U_pred = jax.jit(esn.forecast)(fcast_len=U_test.shape[0], res_state=R[-1]) # feed in the last reservoir state seen in trainingTraceback (most recent call last):
File "/home/mahlau/Programming/tmp/OpenReservoirComputing/tmp.py", line 20, in <module>
U_pred = jax.jit(esn.forecast)(fcast_len=U_test.shape[0], res_state=R[-1]) # feed in the last reservoir state seen in training
File "/home/mahlau/Programming/tmp/OpenReservoirComputing/.venv/lib/python3.13/site-packages/equinox/_module/_module.py", line 554, in __hash__
return hash(
tuple((k, getattr(self, k)) for k in _module_info[type(self)].names_tuple)
)
File "/home/mahlau/Programming/tmp/OpenReservoirComputing/.venv/lib/python3.13/site-packages/equinox/_module/_module.py", line 554, in __hash__
return hash(
tuple((k, getattr(self, k)) for k in _module_info[type(self)].names_tuple)
)
File "/home/mahlau/Programming/tmp/OpenReservoirComputing/.venv/lib/python3.13/site-packages/equinox/_module/_module.py", line 554, in __hash__
return hash(
tuple((k, getattr(self, k)) for k in _module_info[type(self)].names_tuple)
)
TypeError: unhashable type: 'BCOO'
Activity
ymahlau commentedon May 19, 2026
Same happens when trying to jit the training loop:
esn, R = jax.jit(orc.forecaster.train_RCForecaster)(esn, U_train)You should make sure that all common user functions are either fully jittable or it is well-documented which function input and outputs need to static to allow for jitting.
ymahlau commentedon May 19, 2026
openjournals/joss-reviews#10486
Jan-Williams commentedon May 19, 2026
I agree it makes sense to more clearly convey that the functions you mention (and others) are not fully jittable. For ease of use, however, I think it would make more sense to advertise
eqx.filter_jitwhich will handle the static inputs/outputs without issue. The default approach of "If it contains an orc model, useeqx.filter_jit" is more intuitive for users and the advantage of the models being equinox modules. More clearly advertising this would address both of the errors you encountered. What are your thoughts?ymahlau commentedon May 19, 2026
I have had use-cases in the past where it was very important to control which specific parameters are jitted and which arent. So i can think of examples where eqx.filter_jit would not work. I think a more principled approach would be to add
eqx.field(static=true)to all string-based attributes which are never jittable. Then a user can still decide to control manually which parameter is static or not, or just leave it at the default chosen by you.But if you feel like that would require too much re-writing you can also just add more documentation to advertise eqx.filter_jit. I guess for the largest part of users it will make no difference.
Merge pull request #68 from Jan-Williams/improve-jit-documentation
Jan-Williams commentedon May 21, 2026
Addressed in #68
Merge pull request #68 from Jan-Williams/improve-jit-documentation