Make Models Pickleable Again

[alexander-g]

Bugfix for cloudpickling models which also affected Model.save.
Jitted functions cannot be pickled and need to be excluded before pickling.

[cgarciae]

Hey @alexander-g ! I am starting to feel that Model.jit, Model.init_jit and friends should be properties that access a e.g. Model._jit and construct it if its None. We could also add a clear_jit method to clear all jitted functions all at once for stuff like serialization, each should just reconstruct themselves upon access.

I am curious as to why this didn't happened before?

[alexander-g]

I've just checked, it's because of the Jax version. Jitted functions were pickleable up until 0.2.5. Stopped working from 0.2.6 on

[codecov-io]

Codecov Report

Merging #133 (2f356a6) into master (14fba6c) will increase coverage by 0.11%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master     #133      +/-   ##
==========================================
+ Coverage   79.22%   79.34%   +0.11%     
==========================================
  Files         110      110              
  Lines        5190     5210      +20     
==========================================
+ Hits         4112     4134      +22     
+ Misses       1078     1076       -2     

Impacted Files	Coverage Δ
elegy/model/model_base.py	76.42% <100.00%> (+0.51%)	⬆️
elegy/model/model_test.py	100.00% <100.00%> (ø)
elegy/module.py	89.39% <100.00%> (+0.46%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 14fba6c...2f356a6. Read the comment docs.

[cgarciae]

I reviewed the changes a little further and they are good!