Any plans to support keras3?

[Krovatkin]

System information

• TensorFlow version (you are using): 2.12.0
• Are you willing to contribute it (Yes/No): Yes

Motivation
Keras3 is the latest version and it supports multiple backends including Pytorch.

Describe the feature

Supporting keras3 when using different tfmot features including quanitization.
Are there any plans to provide support for keras3? and what needs to be done to support it?

Currently running

import keras
import tensorflow_model_optimization as tfmot

fails with

======================================================================
ERROR: test_qat (__main__.KerasImportTestCase)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "    import_utils/keras3_import_test.runfiles/nuro/learning/py_utils/import_utils/keras3_import_test.py", line 40, in test_qat
    import learning.quantization.quantize_layers.helper as qat_helper
  File "    import_utils/keras3_import_test.runfiles/nuro/learning/quantization/quantize_layers/helper.py", line 4, in <module>
    import tensorflow_model_optimization as tfmot
  File "    import_utils/keras3_import_test.runfiles/nuro/external/pypi__tensorflow_model_optimization_0_8_0/tensorflow_model_optimization/__init__.py", line 86, in <module>
    from tensorflow_model_optimization.python.core.api import clustering
  File "    import_utils/keras3_import_test.runfiles/nuro/external/pypi__tensorflow_model_optimization_0_8_0/tensorflow_model_optimization/python/core/api/__init__.py", line 16, in <module>
    from tensorflow_model_optimization.python.core.api import clustering
  File "    import_utils/keras3_import_test.runfiles/nuro/external/pypi__tensorflow_model_optimization_0_8_0/tensorflow_model_optimization/python/core/api/clustering/__init__.py", line 16, in <module>
    from tensorflow_model_optimization.python.core.api.clustering import keras
  File "    import_utils/keras3_import_test.runfiles/nuro/external/pypi__tensorflow_model_optimization_0_8_0/tensorflow_model_optimization/python/core/api/clustering/keras/__init__.py", line 19, in <module>
    from tensorflow_model_optimization.python.core.clustering.keras.cluster import cluster_scope
  File "    import_utils/keras3_import_test.runfiles/nuro/external/pypi__tensorflow_model_optimization_0_8_0/tensorflow_model_optimization/python/core/clustering/keras/cluster.py", line 22, in <module>
    from tensorflow_model_optimization.python.core.clustering.keras import cluster_wrapper
  File "    import_utils/keras3_import_test.runfiles/nuro/external/pypi__tensorflow_model_optimization_0_8_0/tensorflow_model_optimization/python/core/clustering/keras/cluster_wrapper.py", line 23, in <module>
    from tensorflow_model_optimization.python.core.clustering.keras import clustering_centroids
  File "    import_utils/keras3_import_test.runfiles/nuro/external/pypi__tensorflow_model_optimization_0_8_0/tensorflow_model_optimization/python/core/clustering/keras/clustering_centroids.py", line 22, in <module>
    from tensorflow_model_optimization.python.core.keras.compat import keras
  File "    import_utils/keras3_import_test.runfiles/nuro/external/pypi__tensorflow_model_optimization_0_8_0/tensorflow_model_optimization/python/core/keras/compat.py", line 41, in <module>
    keras = _get_keras_instance()
  File "    import_utils/keras3_import_test.runfiles/nuro/external/pypi__tensorflow_model_optimization_0_8_0/tensorflow_model_optimization/python/core/keras/compat.py", line 33, in _get_keras_instance
    version_fn = getattr(tf.keras, 'version', None)
  File "    import_utils/keras3_import_test.runfiles/nuro/external/org_tensorflow/tensorflow/python/util/lazy_loader.py", line 58, in __getattr__
    module = self._load()
  File "    import_utils/keras3_import_test.runfiles/nuro/external/org_tensorflow/tensorflow/python/util/lazy_loader.py", line 41, in _load
    module = importlib.import_module(self.__name__)
  File "external/local_config_python/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
ModuleNotFoundError: No module named 'keras.api'

[tucan9389]

Hello @Krovatkin, thanks for the issue. Could you provide specific version of Keras 3?

[Krovatkin]

@tucan9389 yes, I'm so terribly sorry for missing yr reply!

We are on

keras=3.0.1
tensorflow=2.12.0 

and python=3.10

[tucan9389]

@Krovatkin

Thanks for the information.

Unfortunately, we don't have a plan to support this library for Keras3 at the moment.

However we're welcome to the contribution. Please contribute if you could fix this error!

cc. @abattery

[apage224]

Just to clarify, TFMOT which is designed specifically with Keras API in mind will not support the latest, default Keras 3 (adopted by TensorFlow) but rather only support the legacy Keras 2? Since TensorFlow 2.16+ targets Keras 3 and TFMOT is still in pre-release, wouldn't it make more sense to migrate TFMOT to Keras 3 as well...

[matpalm]

also note: Keras3 QAT has been "coming soon" ( since Feb 24 ) keras-team/keras#18930