message: 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).
diff: @@ -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
  }
  