New tensorflow version breaks restoring from pickle

[itrharrison]

The recent new version of tensorflow (2.14.0) moves (or removes? I can't actually find the functions in the new version) the tensorflow.python.training.tracking sub-module.

This seems to break restoring networks from pickle files:

(Pdb) tf.__version__
'2.14.0'
(Pdb) filename
'/-----/CosmoPower/CP_paper/CMB/cmb_TT_NN'
(Pdb) pickle.load(open(filename + ".pkl", 'rb'))
*** ModuleNotFoundError: No module named 'tensorflow.python.training.tracking'

I guess the dependency could be pinned tensorflow<2.14.0.

[dpiras]

Hi @itrharrison, thank you for flagging this! Could you confirm restoring from pickle files works correctly with tensorflow<2.14.0?

I will search for a fix, in case I agree we should update the TF dependency as you say.

[dpiras]

I can reproduce the problem with tensorflow==2.14.0, and confirm everything works as expected with e.g. tensorflow==2.13.0. Looking into it now!

[dpiras]

The tracking module seems to be gone in 2.14, and it's not indicated explicitly in the release notes.

I would suggest using 2.13 (or below) for now, and I can open an issue on the TF GitHub page, unless you are already planning to do it @itrharrison?

If they do not plan to fix this, we should specify < 2.14 in the requirements, as you suggested.

[itrharrison]

Issue away! I'm not sure if the pickle restore is their problem or not(?), but it seems wrong that the removal of tracking isn't in the release notes.

[dpiras]

Agreed! Issue opened here, let's see what they say

[dpiras]

Based on tensorflow/tensorflow#62210, they moved the module to tensorflow.python.trackable apparently, and they discouraged the use of pickle. I would say we can add the <2.14.0 in the requirements, and perhaps think of not using pickle in the future.

@itrharrison I can edit the requirements, if you agree?

[itrharrison]

Yup, sounds good!

[dpiras]

Just FYI, I bypassed this problem with a simple script that can be run once to convert the .pkl file to a NumPy dictionary, and then make CosmoPower-JAX usable even with TF>=2.14. If needed, we can probably adapt the same thing for CosmoPower, depending on how you upload your model. Feel free to reopen this issue if you'd need to use TF>=2.14, and we can look into it. This is the CPJ issue.