Bug: Missing function register_module_for_export

[renato-arantes]

What happened?

I'm having an issue with TF v2.16 and tensorflow_hub v0.16, which is throwing an error about a missing function register_module_for_export. If I downgrade to tensorflow_hub v0.15, it complains about a missing estimator.py file, that is deprecated in TF v2.16.

Relevant code

import tensorflow_hub as hub

Relevant log output

import tensorflow_hub as hub
  File "/.../work_dir/venv/lib/python3.11/site-packages/tensorflow_hub/__init__.py", line 111, in <module>
    _ensure_keras_2_importable()
  File "/.../work_dir/venv/lib/python3.11/site-packages/tensorflow_hub/__init__.py", line 100, in _ensure_keras_2_importable
    import tf_keras as keras
  File "/home/.../work_dir/venv/lib/python3.11/site-packages/tf_keras/__init__.py", line 3, in <module>
    from tf_keras import __internal__
  File "/home/.../work_dir/venv/lib/python3.11/site-packages/tf_keras/__internal__/__init__.py", line 3, in <module>
    from tf_keras.__internal__ import backend
  File "/home/.../work_dir/venv/lib/python3.11/site-packages/tf_keras/__internal__/backend/__init__.py", line 3, in <module>
    from tf_keras.src.backend import _initialize_variables as initialize_variables
  File "/home/.../work_dir/venv/lib/python3.11/site-packages/tf_keras/src/__init__.py", line 21, in <module>
    from tf_keras.src import applications
  File "/home/.../work_dir/venv/lib/python3.11/site-packages/tf_keras/src/applications/__init__.py", line 18, in <module>
    from tf_keras.src.applications.convnext import ConvNeXtBase
  File "/home/.../work_dir/venv/lib/python3.11/site-packages/tf_keras/src/applications/convnext.py", line 33, in <module>
    from tf_keras.src.engine import sequential
  File "/home/.../work_dir/venv/lib/python3.11/site-packages/tf_keras/src/engine/sequential.py", line 24, in <module>
    from tf_keras.src.engine import functional
  File "/home/.../work_dir/venv/lib/python3.11/site-packages/tf_keras/src/engine/functional.py", line 33, in <module>
    from tf_keras.src.engine import training as training_lib
  File "/home/.../work_dir/venv/lib/python3.11/site-packages/tf_keras/src/engine/training.py", line 48, in <module>
    from tf_keras.src.saving import saving_api
  File "/home/.../work_dir/venv/lib/python3.11/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 "/home/.../work_dir/venv/lib/python3.11/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 "/home/.../work_dir/venv/lib/python3.11/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'?

tensorflow_hub Version

other (please specify)

TensorFlow Version

other (please specify)

Other libraries

tensorflow_hub==0.16.0 tensorflow==2.16

Python Version

3.x

OS

Linux

[KeijiBranshi]

Hi @renato-arantes. Thanks for posting.

Would you mind sharing the output of pip freeze in your environment?

TensorFlow 2.16 has not been released yet at this time. I'm wondering if you might have a different package installed (eg tf-nightly)

Thanks!

• Brandon from The Kaggle Team

[renato-arantes]

Hi @KeijiBranshi,

Yes, when I say TensorFlow 2.16 I mean tf-nightly.

Thanks!

[KeijiBranshi]

Thanks for the clarification. I'm wondering if this is more a TensorFlow+Keras issue, rather than anything related to tensorflow_hub. See the backwards compatibility documentation here.

Do you have tf-keras-nightly installed as well? iiuc, tf-nightly and tf-keras-nightly are released alongside each other. Keras 3 support doesn't come to TF until 2.16 (ie the current nightly version), so their may be incompatibilities between tf-nightly and the current non-nightly tf-keras (see this discussion).

[ashahba]

I'm seeing a similar issue as well with version 0.16.0 and newer:

$ pip freeze | grep -i 'tensorflow'
tensorflow==2.16.1
tensorflow-hub==0.16.1
tensorflow-io-gcs-filesystem==0.36.0

and when I run this command:

>>> from tensorflow_hub import *
2024-04-04 18:07:24.464029: I external/local_tsl/tsl/cuda/cudart_stub.cc:32] Could not find cuda drivers on your machine, GPU will not be used.
2024-04-04 18:07:24.467517: I external/local_tsl/tsl/cuda/cudart_stub.cc:32] Could not find cuda drivers on your machine, GPU will not be used.
2024-04-04 18:07:24.512078: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-04-04 18:07:25.166988: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: module 'tensorflow_hub' has no attribute 'LatestModuleExporter'

Although according to this: https://github.com/tensorflow/hub/blob/v0.16.1/tensorflow_hub/init.py#L132 It should be packaged as part of the module.

When I looked under site-packages sure enough a lot of modules here https://github.com/tensorflow/hub/tree/v0.16.1/tensorflow_hub are missing from the installed package:

$ ls site-packages/tensorflow_hub/
__init__.py  compressed_module_resolver.py  file_utils.py   module_v2.py  resolver.py  uncompressed_module_resolver.py
__pycache__  config.py                      keras_layer.py  registry.py   tf_utils.py  version.py