Can't use TFMOT version 0.8.0 due to missing dependency

Describe the bug
Can't use latest TFMOT 0.8.0 with the latest TF 2.15.0.post1 on a fresh install because of missing tf_keras dependency. This complicates installations of downstream packages that depend on TFMOT.

System information

TensorFlow version (installed from source or binary): 2.15.0.post1 (installed via pip)

TensorFlow Model Optimization version (installed from source or binary): 0.8.0 (via pip)

Python version: 3.10

Describe the expected behavior

Importing the package should succeed

Describe the current behavior

Importing TFMOT fails

Code to reproduce the issue

import tensorflow_model_optimization

results in:

2024-02-12 22:31:33.815729: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical 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`.
2024-02-12 22:31:33.825475: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
2024-02-12 22:31:33.851921: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-02-12 22:31:33.851950: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-02-12 22:31:33.852716: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2024-02-12 22:31:33.856955: I external/local_tsl/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
2024-02-12 22:31:33.857108: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-02-12 22:31:34.387760: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
WARNING:tensorflow:Your environment has TF_USE_LEGACY_KERAS set to True, but you do not have the tf_keras package installed. You must install it in order to use the legacy tf.keras. Install it via: `pip install tf_keras`
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/vloncar/soft/conda/envs/temp-3.10/lib/python3.10/site-packages/tensorflow_model_optimization/__init__.py", line 86, in <module>
    from tensorflow_model_optimization.python.core.api import clustering
  File "/home/vloncar/soft/conda/envs/temp-3.10/lib/python3.10/site-packages/tensorflow_model_optimization/python/core/api/__init__.py", line 16, in <module>
    from tensorflow_model_optimization.python.core.api import clustering
  File "/home/vloncar/soft/conda/envs/temp-3.10/lib/python3.10/site-packages/tensorflow_model_optimization/python/core/api/clustering/__init__.py", line 16, in <module>
    from tensorflow_model_optimization.python.core.api.clustering import keras
  File "/home/vloncar/soft/conda/envs/temp-3.10/lib/python3.10/site-packages/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 "/home/vloncar/soft/conda/envs/temp-3.10/lib/python3.10/site-packages/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 "/home/vloncar/soft/conda/envs/temp-3.10/lib/python3.10/site-packages/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 "/home/vloncar/soft/conda/envs/temp-3.10/lib/python3.10/site-packages/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 "/home/vloncar/soft/conda/envs/temp-3.10/lib/python3.10/site-packages/tensorflow_model_optimization/python/core/keras/compat.py", line 41, in <module>
    keras = _get_keras_instance()
  File "/home/vloncar/soft/conda/envs/temp-3.10/lib/python3.10/site-packages/tensorflow_model_optimization/python/core/keras/compat.py", line 33, in _get_keras_instance
    version_fn = getattr(tf.keras, 'version', None)
  File "/home/vloncar/soft/conda/envs/temp-3.10/lib/python3.10/site-packages/tensorflow/python/util/lazy_loader.py", line 146, in __getattr__
    self._initialize()
  File "/home/vloncar/soft/conda/envs/temp-3.10/lib/python3.10/site-packages/tensorflow/python/util/lazy_loader.py", line 138, in _initialize
    raise ImportError(  # pylint: disable=raise-missing-from
ImportError: Keras cannot be imported. Check that it is installed.

Additional context
It's clear that the solution is to install tf_keras and indeed everything works as expected afterwards. However, this should happen automatically. Commit 95331880046399a90d7d7533480c159d997aa143 should have fixed that, however the requirement is not present in setup.py file and thus is not installed when using package from PyPI.

When using TF 2.14.1 everything works as expected without additional installations.

Start with setup.py lines 35-41 and compare the declared dependencies with commit 95331880046399a90d7d7533480c159d997aa143. Reproduce a fresh TensorFlow 2.15.0.post1 and TFMOT 0.8.0 installation, then verify that importing tensorflow_model_optimization succeeds without manually installing tf_keras.