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
diff --git a/tensorflow/tools/pip_package/build_pip_package.sh b/tensorflow/tools/pip_package/build_pip_package.sh
index 090ce22..7a07093 100755
--- a/tensorflow/tools/pip_package/build_pip_package.sh
+++ b/tensorflow/tools/pip_package/build_pip_package.sh
@@ -208,16 +208,18 @@
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
}
