Use hermetic Python in TSL and XLA

PiperOrigin-RevId: 634094641

Author: vam-google

WORKSPACE

@@ -7,6 +7,32 @@ workspace(name = "xla")
 # restriction that load() statements need to be at the top of .bzl files.
 # E.g. we can not retrieve a new repository with http_archive and then load()
 # a macro from that repository in the same file.
+
+# Initialize hermetic Python
+load("//third_party/py:python_init_rules.bzl", "python_init_rules")
+
+python_init_rules()
+
+load("//third_party/py:python_init_repositories.bzl", "python_init_repositories")
+
+python_init_repositories(
+    requirements = {
+        "3.11": "//:requirements_lock_3_11.txt",
+    },
+)
+
+load("//third_party/py:python_init_toolchains.bzl", "python_init_toolchains")
+
+python_init_toolchains()
+
+load("//third_party/py:python_init_pip.bzl", "python_init_pip")
+
+python_init_pip()
+
+load("@pypi//:requirements.bzl", "install_deps")
+
+install_deps()
+
 load(":workspace4.bzl", "xla_workspace4")
 
 xla_workspace4()

opensource_only.files

@@ -19,13 +19,17 @@ third_party/llvm_openmp/openmp.bzl:
 third_party/ortools/BUILD:
 third_party/ortools/glpk.BUILD:
 third_party/ortools/ortools.patch:
-third_party/py/non_hermetic/BUILD.tpl:
-third_party/py/non_hermetic/BUILD:
-third_party/py/non_hermetic/README:
-third_party/py/non_hermetic/ml_dtypes/BUILD:
-third_party/py/non_hermetic/ml_dtypes/LICENSE:
-third_party/py/non_hermetic/numpy/BUILD:
-third_party/py/non_hermetic/python_configure.bzl:
+third_party/py/BUILD.tpl:
+third_party/py/BUILD:
+third_party/py/ml_dtypes/BUILD:
+third_party/py/ml_dtypes/LICENSE:
+third_party/py/numpy/BUILD:
+third_party/py/python_configure.bzl:
+third_party/py/python_init_pip.bzl:
+third_party/py/python_init_repositories.bzl:
+third_party/py/python_init_rules.bzl:
+third_party/py/python_init_toolchains.bzl:
+third_party/py/python_repo.bzl:
 third_party/python_runtime/BUILD:
 third_party/repo.bzl:
 third_party/stablehlo/BUILD:

requirements_lock_3_11.txt

@@ -0,0 +1,49 @@
+numpy==1.24.3 \
+    --hash=sha256:0ec87a7084caa559c36e0a2309e4ecb1baa03b687201d0a847c8b0ed476a7187 \
+    --hash=sha256:1a7d6acc2e7524c9955e5c903160aa4ea083736fde7e91276b0e5d98e6332812 \
+    --hash=sha256:202de8f38fc4a45a3eea4b63e2f376e5f2dc64ef0fa692838e31a808520efaf7 \
+    --hash=sha256:210461d87fb02a84ef243cac5e814aad2b7f4be953b32cb53327bb49fd77fbb4 \
+    --hash=sha256:2d926b52ba1367f9acb76b0df6ed21f0b16a1ad87c6720a1121674e5cf63e2b6 \
+    --hash=sha256:352ee00c7f8387b44d19f4cada524586f07379c0d49270f87233983bc5087ca0 \
+    --hash=sha256:35400e6a8d102fd07c71ed7dcadd9eb62ee9a6e84ec159bd48c28235bbb0f8e4 \
+    --hash=sha256:3c1104d3c036fb81ab923f507536daedc718d0ad5a8707c6061cdfd6d184e570 \
+    --hash=sha256:4719d5aefb5189f50887773699eaf94e7d1e02bf36c1a9d353d9f46703758ca4 \
+    --hash=sha256:4749e053a29364d3452c034827102ee100986903263e89884922ef01a0a6fd2f \
+    --hash=sha256:5342cf6aad47943286afa6f1609cad9b4266a05e7f2ec408e2cf7aea7ff69d80 \
+    --hash=sha256:56e48aec79ae238f6e4395886b5eaed058abb7231fb3361ddd7bfdf4eed54289 \
+    --hash=sha256:76e3f4e85fc5d4fd311f6e9b794d0c00e7002ec122be271f2019d63376f1d385 \
+    --hash=sha256:7776ea65423ca6a15255ba1872d82d207bd1e09f6d0894ee4a64678dd2204078 \
+    --hash=sha256:784c6da1a07818491b0ffd63c6bbe5a33deaa0e25a20e1b3ea20cf0e43f8046c \
+    --hash=sha256:8535303847b89aa6b0f00aa1dc62867b5a32923e4d1681a35b5eef2d9591a463 \
+    --hash=sha256:9a7721ec204d3a237225db3e194c25268faf92e19338a35f3a224469cb6039a3 \
+    --hash=sha256:a1d3c026f57ceaad42f8231305d4653d5f05dc6332a730ae5c0bea3513de0950 \
+    --hash=sha256:ab344f1bf21f140adab8e47fdbc7c35a477dc01408791f8ba00d018dd0bc5155 \
+    --hash=sha256:ab5f23af8c16022663a652d3b25dcdc272ac3f83c3af4c02eb8b824e6b3ab9d7 \
+    --hash=sha256:ae8d0be48d1b6ed82588934aaaa179875e7dc4f3d84da18d7eae6eb3f06c242c \
+    --hash=sha256:c91c4afd8abc3908e00a44b2672718905b8611503f7ff87390cc0ac3423fb096 \
+    --hash=sha256:d5036197ecae68d7f491fcdb4df90082b0d4960ca6599ba2659957aafced7c17 \
+    --hash=sha256:d6cc757de514c00b24ae8cf5c876af2a7c3df189028d68c0cb4eaa9cd5afc2bf \
+    --hash=sha256:d933fabd8f6a319e8530d0de4fcc2e6a61917e0b0c271fded460032db42a0fe4 \
+    --hash=sha256:ea8282b9bcfe2b5e7d491d0bf7f3e2da29700cec05b49e64d6246923329f2b02 \
+    --hash=sha256:ecde0f8adef7dfdec993fd54b0f78183051b6580f606111a6d789cd14c61ea0c \
+    --hash=sha256:f21c442fdd2805e91799fbe044a7b999b8571bb0ab0f7850d0cb9641a687092b
+lit==17.0.6 \
+    --hash=sha256:dfa9af9b55fc4509a56be7bf2346f079d7f4a242d583b9f2e0b078fd0abae31b
+ml-dtypes==0.3.2 \
+    --hash=sha256:2c34f2ba9660b21fe1034b608308a01be82bbef2a92fb8199f24dc6bad0d5226 \
+    --hash=sha256:3a17ef2322e60858d93584e9c52a5be7dd6236b056b7fa1ec57f1bb6ba043e33 \
+    --hash=sha256:533059bc5f1764fac071ef54598db358c167c51a718f68f5bb55e3dee79d2967 \
+    --hash=sha256:6604877d567a29bfe7cc02969ae0f2425260e5335505cf5e7fefc3e5465f5655 \
+    --hash=sha256:6b35c4e8ca957c877ac35c79ffa77724ecc3702a1e4b18b08306c03feae597bb \
+    --hash=sha256:763697ab8a88d47443997a7cdf3aac7340049aed45f7521f6b0ec8a0594821fe \
+    --hash=sha256:7a4c3fcbf86fa52d0204f07cfd23947ef05b4ad743a1a988e163caa34a201e5e \
+    --hash=sha256:7afde548890a92b41c0fed3a6c525f1200a5727205f73dc21181a2726571bb53 \
+    --hash=sha256:7ba8e1fafc7fff3e643f453bffa7d082df1678a73286ce8187d3e825e776eb94 \
+    --hash=sha256:91f8783fd1f2c23fd3b9ee5ad66b785dafa58ba3cdb050c4458021fa4d1eb226 \
+    --hash=sha256:93b78f53431c93953f7850bb1b925a17f0ab5d97527e38a7e865b5b4bc5cfc18 \
+    --hash=sha256:961134ea44c7b8ca63eda902a44b58cd8bd670e21d62e255c81fba0a8e70d9b7 \
+    --hash=sha256:b89b194e9501a92d289c1ffd411380baf5daafb9818109a4f49b0a1b6dce4462 \
+    --hash=sha256:c7b3fb3d4f6b39bcd4f6c4b98f406291f0d681a895490ee29a0f95bab850d53c \
+    --hash=sha256:d1a746fe5fb9cd974a91070174258f0be129c592b93f9ce7df6cc336416c3fbd \
+    --hash=sha256:e8505946df1665db01332d885c2020b4cb9e84a8b1241eb4ba69d59591f65855 \
+    --hash=sha256:f47619d978ab1ae7dfdc4052ea97c636c6263e1f19bd1be0e42c346b98d15ff4

third_party/llvm/toolchains.patch

@@ -2,7 +2,7 @@ diff --git a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel b/utils/bazel/llv
 index c43ab727e285..7d848d2dffae 100644
 --- a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel
 +++ b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel
-@@ -30,6 +30,30 @@ exports_files([
+@@ -30,6 +30,36 @@ exports_files([
      "utils/lit/lit.py",
  ])
 
@@ -22,6 +22,12 @@ index c43ab727e285..7d848d2dffae 100644
 +    },
 +)
 +
++py_library(
++    name = "lit_lib",
++    testonly = True,
++    srcs = ["utils/lit/lit.py"] + glob(["utils/lit/lit/**/*.py"]),
++)
++
 +config_setting(
 +    name = "macos_x86_64",
 +    values = {

third_party/py/BUILD

@@ -0,0 +1,40 @@
+load("@python//:defs.bzl", "compile_pip_requirements")
+load("@python_version_repo//:py_version.bzl", "REQUIREMENTS")
+
+compile_pip_requirements(
+    name = "requirements",
+    extra_args = [
+        "--allow-unsafe",
+        "--build-isolation",
+    ],
+    generate_hashes = True,
+    requirements_in = "requirements.in",
+    requirements_txt = REQUIREMENTS,
+)
+
+compile_pip_requirements(
+    name = "requirements_nightly",
+    data = ["test-requirements.txt"],
+    extra_args = [
+        "--allow-unsafe",
+        "--build-isolation",
+        "--extra-index-url=https://pypi.anaconda.org/scientific-python-nightly-wheels/simple",
+        "--pre",
+        "--upgrade",
+    ],
+    generate_hashes = False,
+    requirements_in = "requirements.in",
+    requirements_txt = REQUIREMENTS,
+)
+
+compile_pip_requirements(
+    name = "requirements_dev",
+    extra_args = [
+        "--allow-unsafe",
+        "--build-isolation",
+        "--upgrade",
+    ],
+    generate_hashes = False,
+    requirements_in = "requirements.in",
+    requirements_txt = REQUIREMENTS,
+)

third_party/py/BUILD.tpl

@@ -5,16 +5,17 @@ package(default_visibility = ["//visibility:public"])
 # Point both runtimes to the same python binary to ensure we always
 # use the python binary specified by ./configure.py script.
 load("@bazel_tools//tools/python:toolchain.bzl", "py_runtime_pair")
+load("@python//:defs.bzl", "interpreter")
 
 py_runtime(
     name = "py2_runtime",
-    interpreter_path = "%{PYTHON_BIN_PATH}",
+    interpreter_path = interpreter,
     python_version = "PY2",
 )
 
 py_runtime(
     name = "py3_runtime",
-    interpreter_path = "%{PYTHON_BIN_PATH}",
+    interpreter_path = interpreter,
     python_version = "PY3",
 )
 
@@ -32,27 +33,8 @@ toolchain(
     exec_compatible_with = [%{PLATFORM_CONSTRAINT}],
 )
 
-# To build Python C/C++ extension on Windows, we need to link to python import library pythonXY.lib
-# See https://docs.python.org/3/extending/windows.html
-cc_import(
-    name = "python_lib",
-    interface_library = select({
-        ":windows": ":python_import_lib",
-        # A placeholder for Unix platforms which makes --no_build happy.
-        "//conditions:default": "not-existing.lib",
-    }),
-    system_provided = 1,
-)
-
-cc_library(
-    name = "python_headers",
-    hdrs = [":python_include"],
-    deps = select({
-        ":windows": [":python_lib"],
-        "//conditions:default": [],
-    }),
-    includes = ["python_include"],
-)
+alias(name = "python_headers",
+      actual = "@python//:python_headers")
 
 # This alias is exists for the use of targets in the @llvm-project dependency,
 # which expect a python_headers target called @python_runtime//:headers. We use
@@ -63,18 +45,9 @@ alias(
     actual = ":python_headers",
 )
 
-cc_library(
-    name = "numpy_headers",
-    hdrs = [":numpy_include"],
-    includes = ["numpy_include"],
-)
 
 config_setting(
     name = "windows",
     values = {"cpu": "x64_windows"},
     visibility = ["//visibility:public"],
-)
-
-%{PYTHON_INCLUDE_GENRULE}
-%{NUMPY_INCLUDE_GENRULE}
-%{PYTHON_IMPORT_LIB_GENRULE}
+)

third_party/py/README

@@ -2,15 +2,14 @@ licenses(["restricted"])
 
 package(default_visibility = ["//visibility:public"])
 
-py_library(
+alias(
     name = "numpy",
-    srcs = ["tf_numpy_dummy.py"],
-    srcs_version = "PY3",
+    actual = "@pypi_numpy//:pkg",
 )
 
 alias(
     name = "headers",
-    actual = "@local_config_python//:numpy_headers",
+    actual = "@pypi_numpy//:numpy_headers",
 )
 
 genrule(

third_party/py/numpy/BUILD

@@ -0,0 +1,60 @@
+Copyright (c) 2005-2019, NumPy Developers.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are
+met:
+
+    * Redistributions of source code must retain the above copyright
+       notice, this list of conditions and the following disclaimer.
+
+    * Redistributions in binary form must reproduce the above
+       copyright notice, this list of conditions and the following
+       disclaimer in the documentation and/or other materials provided
+       with the distribution.
+
+    * Neither the name of the NumPy Developers nor the names of any
+       contributors may be used to endorse or promote products derived
+       from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
+
+The NumPy repository and source distributions bundle several libraries that are
+compatibly licensed.  We list these here.
+
+Name: Numpydoc
+Files: doc/sphinxext/numpydoc/*
+License: 2-clause BSD
+  For details, see doc/sphinxext/LICENSE.txt
+
+Name: scipy-sphinx-theme
+Files: doc/scipy-sphinx-theme/*
+License: 3-clause BSD, PSF and Apache 2.0
+  For details, see doc/scipy-sphinx-theme/LICENSE.txt
+
+Name: lapack-lite
+Files: numpy/linalg/lapack_lite/*
+License: 3-clause BSD
+  For details, see numpy/linalg/lapack_lite/LICENSE.txt
+
+Name: tempita
+Files: tools/npy_tempita/*
+License: BSD derived
+  For details, see tools/npy_tempita/license.txt
+
+Name: dragon4
+Files: numpy/core/src/multiarray/dragon4.c
+License: One of a kind
+  For license text, see numpy/core/src/multiarray/dragon4.c

third_party/py/numpy/LICENSE

@@ -390,7 +390,7 @@ tf_module {
   }
   member_method {
     name: "experimental_enable_numpy_behavior"
-    argspec: "args=[\'prefer_float32\'], varargs=None, keywords=None, defaults=[\'False\'], "
+    argspec: "args=[\'prefer_float32\', \'dtype_conversion_mode\'], varargs=None, keywords=None, defaults=[\'False\', \'legacy\'], "
   }
   member_method {
     name: "expm1"