Test pickle support
jheek opened this issue · comments
PR #1475 fixes cloudpickle support but doesn't support add any testing.
Our current test setup seems to import dill indirectly which is incompatible with cloudpickle. Ideally we want to setup testing such that both are tested but don't get in the way of each other. Also we will need to think about how we can test that internal Flax state isn't serialized like was previously the case.