Improve jit documentation

[Jan-Williams]

Opening in response to esn_forecast not jittable.

Two representative issues were raised there:

(1). Code like

 esn, R = orc.forecaster.train_RCForecaster(esn, U_train)

can't be wrapped with jax.jit

 esn, R = jax.jit(orc.forecaster.train_RCForecaster)(esn, U_train)

(2). Code like

 U_pred = esn.forecast(fcast_len=U_test.shape[0], res_state=R[-1])

can't be wrapped with jax.jit

 U_pred = jax.jit(esn.forecast)(fcast_len=U_test.shape[0], res_state=R[-1]) # TypeError: unhashable type: BCOO

(1) has been fixed by explicitly registering string-based attributes as static in ORC layers. esn, R = jax.jit(orc.forecaster.train_RCForecaster)(esn, U_train) (and similar) now run without issue.

(2) is a bit more difficult. JAX sparse BCOO matrices are not hashable, so to my knowledge there is no way to make the above pattern work with jax.jit. eqx.filter_jit, on the other hand, works as a drop in replacement. This is one of the main arguments for the use of eqx.filter_jit over jax.jit and the README.md and docs have been updated to clearly state this preference. jax.jit can still work, but needs something like the following:

@functools.partial(jax.jit, static_argnums=(1,))
def jit_forecast(model, fcast_len, res_state):
    return model.forecast(fcast_len, res_state)

U_pred = jit_forecast(esn, fcast_len, R[-1])
print(f"Forecast shape: {U_pred.shape}")

An additional example notebook has been added to illustrate the use of these different patterns and state the preference for eqx.filter_jit.

@ymahlau Does this address your concerns from esn_forecast not jittable?

[Jan-Williams]

openjournals/joss-reviews#10486

[ymahlau]

@Jan-Williams looks good to me!