From 71313bcb180f2220adee93d964bc4721ad6bf575 Mon Sep 17 00:00:00 2001
From: Anna R <annarev@google.com>
Date: Thu, 19 Mar 2020 13:12:52 -0700
Subject: [PATCH] Place keras API tree at the root of pip package to get
autocomplete working for deeper imports (for e.g. from
tensorflow.keras.losses import KLD).
PiperOrigin-RevId: 301880779
Change-Id: I332e0570750dfdd61712688a4f327f7993019520
---
tensorflow/tools/pip_package/build_pip_package.sh | 12 +++++++-----
1 file changed, 7 insertions(+), 5 deletions(-)
diff --git a/tensorflow/tools/pip_package/build_pip_package.sh b/tensorflow/tools/pip_package/build_pip_package.sh
index 090ce22b718..7a070938045 100755
--- a/tensorflow/tools/pip_package/build_pip_package.sh
+++ b/tensorflow/tools/pip_package/build_pip_package.sh
@@ -208,16 +208,18 @@ function prepare_src() {
rm -f ${TMPDIR}/tensorflow/libtensorflow_framework.so
rm -f ${TMPDIR}/tensorflow/libtensorflow_framework.so.[0-9].*
- # Create a keras/__init__.pyi file so that autocomplete for imports
- # such as `from tensorflow.keras import losses` works.
# TODO(annarev): copy over API files from tensorflow/api/_vN to tensorflow/
# except tensorflow/api/_vN/lite/.
- mkdir ${TMPDIR}/tensorflow/keras/
+
+ # Copy over keras API folder to the root directory
+ # so that autocomplete works as expected for all keras subimports.
if [ -d "${TMPDIR}/tensorflow/_api/v1/" ]
then
- echo "from tensorflow.python.keras.api._v1.keras import *" > ${TMPDIR}/tensorflow/keras/__init__.pyi
+ cp -r ${TMPDIR}/tensorflow/python/keras/api/_v1/keras/ ${TMPDIR}/tensorflow/keras/
+ sed -i'.original' -e 's/.python.keras.api._v1/tensorflow/g' ${TMPDIR}/tensorflow/__init__.py
else
- echo "from tensorflow.python.keras.api._v2.keras import *" > ${TMPDIR}/tensorflow/keras/__init__.pyi
+ cp -r ${TMPDIR}/tensorflow/python/keras/api/_v2/keras/ ${TMPDIR}/tensorflow/keras/
+ sed -i'.original' -e 's/.python.keras.api._v2/tensorflow/g' ${TMPDIR}/tensorflow/__init__.py
fi
}