Fixups for pip_package on Windows

- Use os.path.join instead of hardcoding directory separators when detecting if we're in the tflite_runtime package or not.
- Allow a different file extension for the pywrap library output.
- Have setuptools include any pyd files in the wheel, if they are present.

PiperOrigin-RevId: 334716352
Change-Id: I63760beec3fe7245a3a1e5c12c43b33e9a06c619

tensorflow/lite/python/interpreter.py

@@ -26,7 +26,8 @@ import os
 import numpy as np
 
 # pylint: disable=g-import-not-at-top
-if not os.path.splitext(__file__)[0].endswith('tflite_runtime/interpreter'):
+if not os.path.splitext(__file__)[0].endswith(
+    os.path.join('tflite_runtime', 'interpreter')):
   # This file is part of tensorflow package.
   from tensorflow.lite.python.interpreter_wrapper import _pywrap_tensorflow_interpreter_wrapper as _interpreter_wrapper
   from tensorflow.python.util.tf_export import tf_export as _tf_export

tensorflow/lite/tools/pip_package/README.md

@@ -98,6 +98,12 @@ tensorflow/tools/ci_build/ci_build.sh PI-PYTHON38 \
   tensorflow/lite/tools/pip_package/build_pip_package_with_bazel.sh aarch64
 ```
 
+### Native build for Windows
+
+```sh
+bash tensorflow/lite/tools/pip_package/build_pip_package_with_bazel.sh windows
+```
+
 ## Enable TF OP support (Flex delegate)
 
 If you want to use TF ops with Python API, you need to enable flex support.

tensorflow/lite/tools/pip_package/build_pip_package_with_bazel.sh

@@ -79,10 +79,23 @@ esac
 # include path for Python 3.x builds to work.
 export CROSSTOOL_PYTHON_INCLUDE_PATH
 
+case "${TENSORFLOW_TARGET}" in
+  windows)
+    LIBRARY_EXTENSION=".pyd"
+    ;;
+  *)
+    LIBRARY_EXTENSION=".so"
+    ;;
+esac
+
 bazel build -c opt -s --config=monolithic --config=noaws --config=nogcp --config=nohdfs --config=nonccl \
   ${BAZEL_FLAGS} ${CUSTOM_BAZEL_FLAGS} //tensorflow/lite/python/interpreter_wrapper:_pywrap_tensorflow_interpreter_wrapper
-cp "${TENSORFLOW_DIR}/bazel-bin/tensorflow/lite/python/interpreter_wrapper/_pywrap_tensorflow_interpreter_wrapper.so" \
+cp "${TENSORFLOW_DIR}/bazel-bin/tensorflow/lite/python/interpreter_wrapper/_pywrap_tensorflow_interpreter_wrapper${LIBRARY_EXTENSION}" \
    "${BUILD_DIR}/tflite_runtime"
+# Bazel generates the wrapper library with r-x permissions for user.
+# At least on Windows, we need write permissions to delete the file.
+# Without this, setuptools fails to clean the build directory.
+chmod u+w "${BUILD_DIR}/tflite_runtime/_pywrap_tensorflow_interpreter_wrapper${LIBRARY_EXTENSION}"
 
 # Build python wheel.
 cd "${BUILD_DIR}"

tensorflow/lite/tools/pip_package/setup_with_bazel.py

@@ -63,7 +63,7 @@ setup(
     ],
     packages=find_packages(exclude=[]),
     package_dir={'': '.'},
-    package_data={'': ['*.so']},
+    package_data={'': ['*.so', '*.pyd']},
     install_requires=[
         'numpy >= 1.16.0',
         'pybind11 >= 2.4.3',