Bug: Module Not Found Error 'tensorflow.python.trackable'

[jeipollack]

Launching pytest I encountered a new Tensor Flow Error involving ModuleNotFoundError.

Version:
tensorflow==2.9.1

Traceback

Loading tensorflow-gpu/py3/2.9.1
  Loading requirement: cuda/11.2 nccl/2.9.6-1-cuda cudnn/8.1.1.33-cuda gcc/8.5.0
    openmpi/4.1.1-cuda sox/14.4.2
+ export GITHUB_ENV=true
+ GITHUB_ENV=true
+ cd /gpfswork/rech/ynx/uuu68hq/wf-psf
+ srun pytest -v -s src/wf_psf/tests/metrics_test.py
2023-12-09 21:27:56.240741: I tensorflow/core/util/util.cc:169] oneDNN custom operations are on. You may see slightly different numeri
cal results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `
TF_ENABLE_ONEDNN_OPTS=0`.
ImportError while loading conftest '/gpfsdswork/projects/rech/ynx/uuu68hq/wf-psf/src/wf_psf/tests/conftest.py'.
src/wf_psf/tests/conftest.py:12: in <module>
    from wf_psf.training.train import TrainingParamsHandler
/gpfswork/rech/ynx/uuu68hq/.local/lib/python3.10/site-packages/wf_psf/training/train.py:14: in <module>
    import tensorflow_addons as tfa
/gpfslocalsup/pub/anaconda-py3/2021.05/envs/tensorflow-2.9.1+py3.10/lib/python3.10/site-packages/tensorflow_addons/__init__.py:21: in
<module>
    from tensorflow_addons import activations
/gpfslocalsup/pub/anaconda-py3/2021.05/envs/tensorflow-2.9.1+py3.10/lib/python3.10/site-packages/tensorflow_addons/activations/__init_
_.py:17: in <module>
    from tensorflow_addons.activations.gelu import gelu
/gpfslocalsup/pub/anaconda-py3/2021.05/envs/tensorflow-2.9.1+py3.10/lib/python3.10/site-packages/tensorflow_addons/activations/gelu.py
:19: in <module>
    from tensorflow_addons.utils.types import TensorLike
/gpfslocalsup/pub/anaconda-py3/2021.05/envs/tensorflow-2.9.1+py3.10/lib/python3.10/site-packages/tensorflow_addons/utils/types.py:45:
in <module>
    Initializer = Union[None, dict, str, Callable, tf.keras.initializers.Initializer]
/gpfslocalsup/pub/anaconda-py3/2021.05/envs/tensorflow-2.9.1+py3.10/lib/python3.10/site-packages/tensorflow/python/util/lazy_loader.py
:58: in __getattr__
    module = self._load()
/gpfslocalsup/pub/anaconda-py3/2021.05/envs/tensorflow-2.9.1+py3.10/lib/python3.10/site-packages/tensorflow/python/util/lazy_loader.py
:41: in _load
    module = importlib.import_module(self.__name__)
/gpfswork/rech/ynx/uuu68hq/.local/lib/python3.10/site-packages/keras/__init__.py:21: in <module>
    from keras import models
/gpfswork/rech/ynx/uuu68hq/.local/lib/python3.10/site-packages/keras/models/__init__.py:18: in <module>
    from keras.engine.functional import Functional
/gpfswork/rech/ynx/uuu68hq/.local/lib/python3.10/site-packages/keras/engine/functional.py:27: in <module>
    from keras.dtensor import layout_map as layout_map_lib
/gpfswork/rech/ynx/uuu68hq/.local/lib/python3.10/site-packages/keras/dtensor/layout_map.py:25: in <module>
    from keras.dtensor import lazy_variable
/gpfswork/rech/ynx/uuu68hq/.local/lib/python3.10/site-packages/keras/dtensor/lazy_variable.py:26: in <module>
    from tensorflow.python.trackable import base as trackable
E   ModuleNotFoundError: No module named 'tensorflow.python.trackable'

[jeipollack]

This solution also solves #93.