Fix #1272: tensorflow-model-optimization fails to import in official ...

[JiwaniZakir]

Closes #1272

Motivation

In the official tensorflow/tensorflow:2.15.0 Docker image, TF_USE_LEGACY_KERAS=1 is preset, which causes tf.keras.version() to return a Keras 3.x string. compat._get_keras_instance() then attempts import tf_keras, which is not present in that image, raising an ImportError that propagates and makes the entire tensorflow_model_optimization package unimportable.

Changes

tensorflow_model_optimization/python/core/keras/compat.py

In _get_keras_instance(), the unconditional import tf_keras as keras_internal (previously line 35) is now wrapped in a try/except ImportError. When tf_keras is not installed, the function falls back to keras_internal = tf.keras — the Keras 2 instance bundled with TensorFlow — rather than raising. The else branch for non-Keras-3 environments is unchanged.

tensorflow_model_optimization/python/core/keras/metrics_test.py

Adds CompatGetKerasInstanceTest.test_fallback_to_tf_keras_when_tf_keras_not_installed, which:

• Removes tf_keras from sys.modules to simulate it being absent.
• Patches tf.keras.version to return '3.0.0' (triggering the if branch).
• Patches sys.modules['tf_keras'] = None to force the ImportError on import.
• Asserts that compat._get_keras_instance() returns tf.keras rather than raising.

Testing

The new unit test covers the fallback path directly and can be run with:

python -m pytest tensorflow_model_optimization/python/core/keras/metrics_test.py -v -k CompatGetKerasInstanceTest

Manual verification in tensorflow/tensorflow:2.15.0:

docker run --rm tensorflow/tensorflow:2.15.0 bash -c \
  "pip install -q tensorflow-model-optimization && \
   python -c 'import tensorflow_model_optimization as tfmot; print(tfmot.__version__)'"

With the fix applied this completes without error; before the fix it exits with ImportError: Keras cannot be imported. Check that it is installed.

---

This PR was created with AI assistance (Claude). The changes were reviewed by quality gates and a critic model before submission.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.