Update keras base_dir for new source location.

PiperOrigin-RevId: 380052162
Change-Id: Ic142aeb90ccc57a7a51828de92782729c157e847

tensorflow/tools/docs/base_dir.py

@@ -26,12 +26,25 @@
 import tensorflow as tf
 import tensorflow_estimator
 
+try:
+  import keras  # pylint: disable=g-import-not-at-top
+except ImportError:
+  pass
+
 
 def get_base_dirs_and_prefixes(code_url_prefix):
   """Returns the base_dirs and code_prefixes for OSS TensorFlow api gen."""
   base_dir = path.dirname(tf.__file__)
 
-  if distutils.version.LooseVersion(tf.__version__) >= "2.2":
+  if distutils.version.LooseVersion(tf.__version__) >= "2.6":
+    base_dirs = [
+        base_dir,
+        path.dirname(keras.__file__),
+        path.dirname(keras_preprocessing.__file__),
+        path.dirname(tensorboard.__file__),
+        path.dirname(tensorflow_estimator.__file__),
+    ]
+  elif distutils.version.LooseVersion(tf.__version__) >= "2.2":
     base_dirs = [
         base_dir,
         path.dirname(keras_preprocessing.__file__),
@@ -46,11 +59,20 @@ def get_base_dirs_and_prefixes(code_url_prefix):
         path.dirname(tensorflow_estimator.__file__),
     ]
 
-  code_url_prefixes = (
-      code_url_prefix,
-      "https://github.com/keras-team/keras-preprocessing/tree/master/keras_preprocessing",
-      "https://github.com/tensorflow/tensorboard/tree/master/tensorboard",
-      "https://github.com/tensorflow/estimator/tree/master/tensorflow_estimator",
-  )
+  if distutils.version.LooseVersion(tf.__version__) >= "2.6":
+    code_url_prefixes = (
+        code_url_prefix,
+        "https://github.com/keras-team/keras/tree/master/keras",
+        "https://github.com/keras-team/keras-preprocessing/tree/master/keras_preprocessing",
+        "https://github.com/tensorflow/tensorboard/tree/master/tensorboard",
+        "https://github.com/tensorflow/estimator/tree/master/tensorflow_estimator",
+    )
+  else:
+    code_url_prefixes = (
+        code_url_prefix,
+        "https://github.com/keras-team/keras-preprocessing/tree/master/keras_preprocessing",
+        "https://github.com/tensorflow/tensorboard/tree/master/tensorboard",
+        "https://github.com/tensorflow/estimator/tree/master/tensorflow_estimator",
+    )
 
   return base_dirs, code_url_prefixes