deepcopy broken for new FlatMapping about dm-haiku HOT 2 CLOSED

Thanks for noticing, this is mostly @LenaMartens hard work!

We're working on a patch to get all our frozendict tests passing for FlatMapping and this includes covering copy/deepcopy/pickle etc. It's basically there, a missing puzzle piece was google-deepmind/tree@f18a0f6 and I hope once a new dm-tree release is pushed (I think we can do this today) it will be good to go 😄

For pickle support I was planning to go with:

  def __reduce__(self):
    return FlatMapping.from_mapping, (self.to_mapping(),)

I think this is nice since it is forward and backwards compatible (e.g. a FlatMapping saved by Haiku version N would be loadable by Haiku version N-1 and N+1) assuming the pickled representation of dict does not itself change (e.g. if you keep your Python version the same this will be the case and typically between minor versions the pickle format for builtins does not change).

I'll tag this issue with the changes so you're kept up to date.

from dm-haiku.

Awesome, thanks!

from dm-haiku.