Add import version check for tf keras module. Make sure TFMA depend on keras 2 to fix the metrics module mismatch when using TFMA with keras metrics.

PiperOrigin-RevId: 604720506

Author: zhouhao138

tf_keras/__init__.py

@@ -24,6 +24,7 @@
 from tf_keras import losses
 from tf_keras import metrics
 from tf_keras import models
+from tf_keras import optimizers
 from tf_keras.engine.input_layer import Input
 from tf_keras.engine.sequential import Sequential
 from tf_keras.engine.training import Model