Importing `tf_keras` to use Keras 2 in TensorFlow 2.16 fails

[MaxiBoether]

System information.

• Have I written custom code (as opposed to using a stock example script provided in Keras): No.
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macOS 14.2.1, Python 3.12.2 inside a conda environment
• TensorFlow installed from (source or binary): binary (pip)
• TensorFlow version (use command below): 2.16
• Python version: 3.12.2
• Bazel version (if compiling from source): -
• GPU model and memory: -
  Describe the problem.

huggingface transformers currently relies on Keras 2. They mention that install the tf-keras package can make Keras 2 APIs available in TF 2.16. The getting started page mentions something similar. In a clean environment, I install the following packages:

pip install tensorflow tf_keras

Importing Keras 3.0 works as expected:

python -c "import tensorflow.keras; print(tensorflow.keras.version())"  
3.0.5

However, importing tf_keras fails with some error:

 python -c "import tf_keras"          
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/__init__.py", line 3, in <module>
    from tf_keras import __internal__
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/__internal__/__init__.py", line 3, in <module>
    from tf_keras.__internal__ import backend
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/__internal__/backend/__init__.py", line 3, in <module>
    from tf_keras.src.backend import _initialize_variables as initialize_variables
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/__init__.py", line 21, in <module>
    from tf_keras.src import applications
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/applications/__init__.py", line 18, in <module>
    from tf_keras.src.applications.convnext import ConvNeXtBase
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/applications/convnext.py", line 33, in <module>
    from tf_keras.src.engine import sequential
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/engine/sequential.py", line 24, in <module>
    from tf_keras.src.engine import functional
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/engine/functional.py", line 33, in <module>
    from tf_keras.src.engine import training as training_lib
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/engine/training.py", line 48, in <module>
    from tf_keras.src.saving import saving_api
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/saving/saving_api.py", line 25, in <module>
    from tf_keras.src.saving.legacy import save as legacy_sm_saving_lib
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/saving/legacy/save.py", line 27, in <module>
    from tf_keras.src.saving.legacy.saved_model import load_context
  File "/Users/mboether/mambaforge/envs/compile2/lib/python3.12/site-packages/tf_keras/src/saving/legacy/saved_model/load_context.py", line 68, in <module>
    tf.__internal__.register_load_context_function(in_load_context)
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'tensorflow._api.v2.compat.v2.__internal__' has no attribute 'register_load_context_function'. Did you mean: 'register_call_context_function'?

I can trigger the same error by forcing Tensorflow to use Keras 2 instead using TF_USE_LEGACY_KERAS=1 and importing tensorflow.keras. I do not think this is supposed to happen. From my understanding, installing tf_keras should enable us to use keras 2 in tf 2.16.

[tilakrayal]

@MaxiBoether,
Could you please try to replace it from tensorflow import keras with import tf_keras as keras. Kindly find the gist for the reference.

Thank you!

[MaxiBoether]

Thanks for the swift response. This leads to the same error as outlined. Note that I tried two things: import tensorflow.keras (which is equivalent to from tensorflow import keras) which imports keras 3 as expected, but importing import tf_keras (it shouldn't matter whether I add as keras as requested since that's just the imported module name) fails. My goal is to import the legacy Keras 2 in a tf2.16 environment

[MaxiBoether]

Okay, I know what the issue is: By default (at least on my mac in a clean conda env), pip install tensorflow installs tf2.16, and pip install tf_keras installs tf_keras==2.15. On colab, pip install tensorflow defaults to v2.15. Since I just ran pip install tensorflow tf_keras, we run into this issue.

The version mismatch (if enforced on colab) leads to this error.

[robertbcalhoun]

I'm having the same issue; attempting import tensorflow_hub (which imports tf_keras) fails with this error:

AttributeError: module 'tensorflow._api.v2.compat.v2.__internal__' 
has no attribute 'register_load_context_function'.
   Did you mean: 'register_call_context_function'?

This seems to be the result of some in-process refactoring of tf_keras/src/saving/legacy/saved_model that has since been completed, so presumably a later release will fix.

At the moment, the default installs of tensorflow and tensorflow_hub install tensorflow 2.16.0rc0 and tf-keras 2.15:

$ # start with new venv
$ python -m pip freeze | grep -e tensor -e keras
$ python -m pip install tensorflow keras tf-keras --no-cache-dir > /dev/null
$ python -m pip freeze | grep -e tensor -e keras
keras==3.0.5
tensorboard==2.16.2
tensorboard-data-server==0.7.2
tensorflow==2.16.0rc0
tf-keras==2.15.0

I've been trying to work around the incompatibility by reverting to a prior tensorflow release. I expect to be able to install out-of-date versions via pip like this:

$ python -m pip index versions tf-keras
WARNING: pip index is currently an experimental command.
 It may be removed/changed in a future release without prior warning.
tf-keras (2.15.0)
Available versions: 2.15.0, 2.14.1
$ python -m pip install tf-keras==2.14.1
Collecting tf-keras==2.14.1
  Using cached tf_keras-2.14.1-py3-none-any.whl.metadata (1.6 kB)
Using cached tf_keras-2.14.1-py3-none-any.whl (1.7 MB)
Installing collected packages: tf-keras
Successfully installed tf-keras-2.14.1

However, this approach is not working with tensorflow. I don't see anything obviously broken at PyPI, but I've not been able to install prior tensorflow releases via pip.

$ python -m pip install tensorflow==2.15.0
ERROR: Could not find a version that satisfies the requirement tensorflow==2.15.0 (from versions: 2.16.0rc0)
ERROR: No matching distribution found for tensorflow==2.15.0
$ python -m pip index versions tensorflow
WARNING: pip index is currently an experimental command. 
It may be removed/changed in a future release without prior warning.
ERROR: No matching distribution found for tensorflow

[robertbcalhoun]

It appears 2.16.0rc0 is a pre-release version of tensorflow. Running the same commands on a different host (ubuntu under wsl) gives me 2.15.0.post1, which is what I would expect based on PyPI:

$ python -m pip index versions tensorflow
tensorflow (2.15.0.post1)
Available versions: 2.15.0.post1, 2.15.0, 2.14.1, 2.14.0, 2.13.1, 2.13.0, 2.12.1, 2.12.0, 2.11.1, 2.11.0, 2.10.1, 2.10.0, 2.9.3, 2.9.2, 2.9.1, 2.9.0, 2.8.4, 2.8.3, 2.8.2, 2.8.1, 2.8.0
  INSTALLED: 2.15.0.post1
  LATEST:    2.15.0.post1

Overall:

$ python -m pip freeze | grep -e tensor -e keras
keras==2.15.0
tensorboard==2.15.2
tensorboard-data-server==0.7.2
tensorflow==2.15.0.post1
tensorflow-estimator==2.15.0
tensorflow-hub==0.16.1
tensorflow-io-gcs-filesystem==0.36.0
tf-keras==2.15.0

With tensorflow 2.15.0.post1, tf-keras loads correctly under tensorflow-hub. I must have somehow misconfigured pip to pull pre-release versions on the other machine. Time for a rebuild...

[sampathweb]

When you use TensorFlow 2.16 with RC0 and testing Keras 2 path, please install pip install tf-keras~=2.16.0rc0 (will install RC2). If you use pip install tf-keras instead, it will install tf-keras==2.15 which is not compatible with TensorFlow 2.16.

The core of the issue is, in tf-keras 2.15 we didn't add a dependency to TensorFlow 2.15, so pip doesn't enforce that version compatibility. We fixed it in tf-keras 2.16, but may be we need to add a patch release for tf-keras 2.15 to avoid this.

@robertbcalhoun - Can you try this and confirm?

[github-actions]

This issue is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.

[MaxiBoether]

Just trying to avoid this issue becoming stale :-)