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

@@ -20,6 +20,7 @@
 
 from tf_keras import applications
 from tf_keras import distribute
+from tf_keras import layers
 from tf_keras import losses
 from tf_keras import metrics
 from tf_keras import models