From 5720ab7845de0b2a2e2f5fdf547d2515d39a20b9 Mon Sep 17 00:00:00 2001 From: Vadym Matsishevskyi Date: Thu, 10 Oct 2024 00:02:57 -0700 Subject: [PATCH] Introduce pywrap bazel rules and migrate Tensorflow to it The gist of this change (the new rules implementation) is contained within `rules_pywrap` folder. The rules are generic and not tensorflow-specific 1) (internal-specific) 2) (internal-specific) 3) It provides same linking strategy of final artifacts on all 3 supported platforms (no major differences between Linux, Mac and Windows). 4) It makes it possible to abandon usage of header-only targets to prevent ODR violations. Simply speaking you can now depend on generated protobuf message classes normally, without need to worry how that is linked afterwards. 5) The current version is backward-compatible and unless explicitly enabled is a no-op. To enable the new rules pass `--repo_env=USE_PYWRAP_RULES=True` flag to build/test command. 6) The `if_pywrap` construct is temporary and will be removed once full migration is completed. Currently if_pywrap is mainly used to pass normal dependencies (instead of header-only). The header-only stuff is kept for backward compatibility and smoother migration but will be eventually removed. 7) This CL migrates TF and the most problematic among all google ML repositories. Once TF is sabilized the other repositories, such as JAX and XLA will be migrated too (which should be way easier than migrating TF anyways) PiperOrigin-RevId: 684324990 --- WORKSPACE | 8 +- tensorflow/BUILD | 13 +- tensorflow/c/BUILD | 7 +- tensorflow/c/experimental/gradients/BUILD | 4 +- tensorflow/compiler/aot/tfcompile.bzl | 1 + tensorflow/compiler/mlir/lite/python/BUILD | 6 +- .../mlir/quantization/stablehlo/python/BUILD | 3 + .../mlir/quantization/tensorflow/BUILD | 1 + .../mlir/quantization/tensorflow/python/BUILD | 8 +- tensorflow/compiler/mlir/stablehlo/BUILD | 5 + tensorflow/compiler/mlir/tfr/BUILD | 3 + tensorflow/compiler/mlir/tfr/build_defs.bzl | 16 +- tensorflow/compiler/tf2xla/BUILD | 5 +- tensorflow/core/BUILD | 13 +- tensorflow/core/framework/BUILD | 4 + .../core/platform/build_config_root.bzl | 2 + .../platform/build_config_root.default.bzl | 6 + tensorflow/dtensor/python/tests/BUILD | 3 + tensorflow/lite/python/metrics/BUILD | 3 +- tensorflow/python/BUILD | 251 ++++- tensorflow/python/client/BUILD | 48 +- tensorflow/python/framework/BUILD | 89 +- .../python/framework/experimental/BUILD | 33 +- tensorflow/python/grappler/BUILD | 52 +- tensorflow/python/kernel_tests/proto/BUILD | 33 +- tensorflow/python/profiler/internal/BUILD | 2 + .../protobuf_inline_symbols_enforcer.cc | 91 ++ tensorflow/python/saved_model/BUILD | 1 + tensorflow/python/tools/BUILD | 12 +- tensorflow/python/tools/tools.bzl | 1 + tensorflow/python/tpu/BUILD | 11 +- tensorflow/python/util/BUILD | 37 +- tensorflow/tensorflow.bzl | 171 +++- tensorflow/tensorflow.default.bzl | 19 + tensorflow/tools/common/BUILD | 25 +- tensorflow/workspace2.bzl | 5 +- third_party/protobuf/protobuf.patch | 30 +- third_party/py/python_repo.bzl | 10 + third_party/pybind11_protobuf/protobuf.patch | 20 + .../py/non_hermetic/ml_dtypes/BUILD | 0 .../py/non_hermetic/ml_dtypes/LICENSE | 202 ---- .../py/non_hermetic/ml_dtypes/ml_dtypes.BUILD | 64 -- .../ml_dtypes/ml_dtypes.tests.BUILD | 71 -- .../py/non_hermetic/ml_dtypes/workspace.bzl | 22 - .../third_party/py/non_hermetic/numpy/BUILD | 21 - .../py/non_hermetic/numpy/README.md | 4 - .../py/non_hermetic/numpy/tf_numpy_api/BUILD | 12 - ...ensorflow.experimental.numpy.ndarray.pbtxt | 51 - .../tensorflow.experimental.numpy.pbtxt | 919 ------------------ ...tensorflow.experimental.numpy.random.pbtxt | 35 - .../xla/third_party/py/python_repo.bzl | 10 + .../tsl/third_party/protobuf/protobuf.patch | 30 +- .../py/non_hermetic/ml_dtypes/BUILD | 0 .../py/non_hermetic/ml_dtypes/LICENSE | 202 ---- .../py/non_hermetic/ml_dtypes/ml_dtypes.BUILD | 64 -- .../ml_dtypes/ml_dtypes.tests.BUILD | 71 -- .../py/non_hermetic/ml_dtypes/workspace.bzl | 22 - .../third_party/py/non_hermetic/numpy/BUILD | 21 - .../py/non_hermetic/numpy/README.md | 4 - .../py/non_hermetic/numpy/tf_numpy_api/BUILD | 12 - ...ensorflow.experimental.numpy.ndarray.pbtxt | 51 - .../tensorflow.experimental.numpy.pbtxt | 919 ------------------ ...tensorflow.experimental.numpy.random.pbtxt | 35 - .../tsl/third_party/py/python_repo.bzl | 10 + .../tsl/third_party/py/rules_pywrap/BUILD | 16 + .../py/rules_pywrap/pybind_extension.py.tpl | 49 + .../third_party/py/rules_pywrap/pywrap.bzl | 20 + .../py/rules_pywrap/pywrap.default.bzl | 146 +++ .../py/rules_pywrap/pywrap.impl.bzl | 731 ++++++++++++++ .../tsl/tsl/platform/build_config_root.bzl | 2 + .../tsl/tsl/platform/default/BUILD | 1 + .../tsl/tsl/platform/default/build_config.bzl | 88 +- .../platform/default/build_config_root.bzl | 13 + third_party/xla/xla/stream_executor/BUILD | 1 + third_party/xla/xla/tsl/BUILD | 1 + third_party/xla/xla/tsl/tsl.bzl | 48 +- third_party/xla/xla/tsl/tsl.default.bzl | 2 +- 77 files changed, 1994 insertions(+), 3028 deletions(-) create mode 100644 tensorflow/python/protobuf_inline_symbols_enforcer.cc create mode 100644 third_party/pybind11_protobuf/protobuf.patch delete mode 100644 third_party/xla/third_party/py/non_hermetic/ml_dtypes/BUILD delete mode 100644 third_party/xla/third_party/py/non_hermetic/ml_dtypes/LICENSE delete mode 100644 third_party/xla/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.BUILD delete mode 100644 third_party/xla/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.tests.BUILD delete mode 100644 third_party/xla/third_party/py/non_hermetic/ml_dtypes/workspace.bzl delete mode 100644 third_party/xla/third_party/py/non_hermetic/numpy/BUILD delete mode 100644 third_party/xla/third_party/py/non_hermetic/numpy/README.md delete mode 100644 third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/BUILD delete mode 100644 third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.ndarray.pbtxt delete mode 100644 third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.pbtxt delete mode 100644 third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.random.pbtxt delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/BUILD delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/LICENSE delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.BUILD delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.tests.BUILD delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/workspace.bzl delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/BUILD delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/README.md delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/BUILD delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.ndarray.pbtxt delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.pbtxt delete mode 100644 third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.random.pbtxt create mode 100644 third_party/xla/third_party/tsl/third_party/py/rules_pywrap/BUILD create mode 100644 third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pybind_extension.py.tpl create mode 100644 third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pywrap.bzl create mode 100644 third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pywrap.default.bzl create mode 100644 third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pywrap.impl.bzl diff --git a/WORKSPACE b/WORKSPACE index 73cb4ee89c9..0171c60db38 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -24,11 +24,11 @@ load("@//tensorflow:workspace3.bzl", "tf_workspace3") tf_workspace3() # Initialize hermetic Python -load("@local_xla//third_party/py:python_init_rules.bzl", "python_init_rules") +load("@local_tsl//third_party/py:python_init_rules.bzl", "python_init_rules") python_init_rules() -load("@local_xla//third_party/py:python_init_repositories.bzl", "python_init_repositories") +load("@local_tsl//third_party/py:python_init_repositories.bzl", "python_init_repositories") python_init_repositories( default_python_version = "system", @@ -46,11 +46,11 @@ python_init_repositories( }, ) -load("@local_xla//third_party/py:python_init_toolchains.bzl", "python_init_toolchains") +load("@local_tsl//third_party/py:python_init_toolchains.bzl", "python_init_toolchains") python_init_toolchains() -load("@local_xla//third_party/py:python_init_pip.bzl", "python_init_pip") +load("@local_tsl//third_party/py:python_init_pip.bzl", "python_init_pip") python_init_pip() diff --git a/tensorflow/BUILD b/tensorflow/BUILD index c3e28b2f55d..40c24a54813 100644 --- a/tensorflow/BUILD +++ b/tensorflow/BUILD @@ -13,7 +13,7 @@ load( "if_google", "if_oss", "if_xla_available", - "tf_cc_shared_object", + "pywrap_aware_tf_cc_shared_object", "tf_custom_op_library_additional_deps_impl", "tf_monitoring_python_deps", "tf_native_cc_binary", @@ -1085,20 +1085,13 @@ bzl_library( "@local_config_cuda//cuda:build_defs_bzl", "@local_config_rocm//rocm:build_defs_bzl", "@local_config_tensorrt//:build_defs_bzl", - "@local_tsl//tsl/platform/default:cuda_build_defs_bzl", + "@local_tsl//third_party/py/rules_pywrap:pywrap_bzl", "@local_xla//xla/tsl:tsl_bzl", "@local_xla//xla/tsl/mkl:build_defs_bzl", "@rules_java//java:rules", ], ) -bzl_library( - name = "tensorflow_default_bzl", - srcs = ["tensorflow.default.bzl"], - visibility = ["//visibility:public"], - deps = [":tensorflow_bzl"], -) - # TODO(jakeharmon8): Remove these in favor of tsl:grpc # copybara:comment_begin(oss-only) cc_library( @@ -1232,7 +1225,7 @@ cc_import( # an "-exported_symbols_list" command. -z defs disallows undefined # symbols in object files. -tf_cc_shared_object( +pywrap_aware_tf_cc_shared_object( name = "tensorflow", linkopts = select({ "//tensorflow:macos": [ diff --git a/tensorflow/c/BUILD b/tensorflow/c/BUILD index 4380e1cc701..b5729ea7a3d 100644 --- a/tensorflow/c/BUILD +++ b/tensorflow/c/BUILD @@ -3,6 +3,7 @@ load("@bazel_skylib//lib:selects.bzl", "selects") load("@local_config_tensorrt//:build_defs.bzl", "if_tensorrt") +load("@local_tsl//tsl/platform:build_config_root.bzl", "if_pywrap") load("@local_tsl//tsl/platform:rules_cc.bzl", "cc_library") load( "//tensorflow:tensorflow.bzl", @@ -932,6 +933,7 @@ tf_cuda_cc_test( ":test_op1.so", "//tensorflow/cc/saved_model:saved_model_half_plus_two", ], + extra_copts = if_pywrap(["-DTENSORFLOW_NO_SHARED_OBJECTS"]), linkopts = select({ "//tensorflow:macos": ["-headerpad_max_install_names"], "//conditions:default": [], @@ -1131,7 +1133,10 @@ tf_cuda_library( name = "python_api", srcs = ["python_api.cc"], hdrs = ["python_api.h"], - visibility = ["//tensorflow/python:__pkg__"], + visibility = [ + "//tensorflow:internal", + "//tensorflow/python:__pkg__", + ], deps = [ ":c_api", ":c_api_internal", diff --git a/tensorflow/c/experimental/gradients/BUILD b/tensorflow/c/experimental/gradients/BUILD index a3fa49fffa3..92c87960986 100644 --- a/tensorflow/c/experimental/gradients/BUILD +++ b/tensorflow/c/experimental/gradients/BUILD @@ -55,7 +55,9 @@ cc_library( hdrs = [ "nn_grad.h", ], - visibility = ["//visibility:private"], + visibility = [ + "//tensorflow/python/framework/experimental:__pkg__", + ], deps = [ "//tensorflow/c/eager:abstract_tensor_handle", "//tensorflow/c/eager:gradients_internal", diff --git a/tensorflow/compiler/aot/tfcompile.bzl b/tensorflow/compiler/aot/tfcompile.bzl index 99c8541c554..82fdb603138 100644 --- a/tensorflow/compiler/aot/tfcompile.bzl +++ b/tensorflow/compiler/aot/tfcompile.bzl @@ -212,6 +212,7 @@ def _tf_library( ] + freeze_saver_srcs, outs = [freeze_file], cmd = ( + "PYWRAP_TARGET='//tensorflow/python:_pywrap_tensorflow' " + "CUDA_VISIBLE_DEVICES='' " + "$(location " + "//tensorflow/python/tools:freeze_graph)" + diff --git a/tensorflow/compiler/mlir/lite/python/BUILD b/tensorflow/compiler/mlir/lite/python/BUILD index 5f2bcd5f8b5..7d6eb5cef48 100644 --- a/tensorflow/compiler/mlir/lite/python/BUILD +++ b/tensorflow/compiler/mlir/lite/python/BUILD @@ -1,3 +1,4 @@ +load("@local_tsl//tsl/platform:build_config_root.bzl", "if_pywrap") load("//tensorflow:strict.default.bzl", "py_strict_library") load("//tensorflow:tensorflow.default.bzl", "tf_python_pybind_extension") load("//tensorflow/core/platform:rules_cc.bzl", "cc_library") @@ -239,10 +240,13 @@ tf_python_pybind_extension( pytype_srcs = [ "_pywrap_converter_api.pyi", ], + visibility = [ + "//tensorflow/python:__pkg__", + ], deps = [ "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib", "//tensorflow/python/lib/core:pybind11_lib", "//third_party/python_runtime:headers", "@pybind11", - ], + ] + if_pywrap([":converter_python_api"]), ) diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD b/tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD index 5ca03bfc209..a022cbc9e05 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD +++ b/tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD @@ -151,6 +151,9 @@ tf_python_pybind_extension( name = "pywrap_quantization", srcs = ["pywrap_quantization.cc"], pytype_srcs = ["pywrap_quantization.pyi"], + visibility = [ + "//tensorflow/python:__pkg__", + ], # Each dependency MUST be either header-only or exclusive. deps = [ ":pywrap_quantization_lib_header_only", diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/BUILD b/tensorflow/compiler/mlir/quantization/tensorflow/BUILD index e5696e4db86..b9b5aded172 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/BUILD +++ b/tensorflow/compiler/mlir/quantization/tensorflow/BUILD @@ -575,6 +575,7 @@ tf_proto_library( ":internal_visibility_allowlist_package", # To be visible from `lib_internal_impl`. "//tensorflow/core:__pkg__", + "//tensorflow/python:__pkg__", ], ) diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/python/BUILD b/tensorflow/compiler/mlir/quantization/tensorflow/python/BUILD index 41e2f04651a..1f118305309 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/python/BUILD +++ b/tensorflow/compiler/mlir/quantization/tensorflow/python/BUILD @@ -1,3 +1,4 @@ +load("@local_tsl//tsl/platform:build_config_root.bzl", "if_pywrap") load("//tensorflow:pytype.default.bzl", "pytype_strict_library") load( "//tensorflow:tensorflow.default.bzl", @@ -203,6 +204,7 @@ tf_python_pybind_extension( pytype_srcs = ["pywrap_function_lib.pyi"], visibility = [ "__subpackages__", + "//tensorflow/python:__pkg__", "//tensorflow/tools/pip_package:__subpackages__", ], deps = [ @@ -230,7 +232,6 @@ tf_python_pybind_extension( # All deps must be header-only. deps = [ ":py_function_lib", - ":quantize_model_cc", ":type_casters", "//tensorflow/compiler/mlir/quantization/tensorflow:exported_model_proto_cc", "//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc", @@ -246,7 +247,10 @@ tf_python_pybind_extension( "@pybind11_abseil//pybind11_abseil:import_status_module", "@pybind11_abseil//pybind11_abseil:status_casters", "@pybind11_protobuf//pybind11_protobuf:native_proto_caster", - ], + ] + if_pywrap( + [":quantize_model_cc_impl"], + [":quantize_model_cc"], + ), ) tf_py_strict_test( diff --git a/tensorflow/compiler/mlir/stablehlo/BUILD b/tensorflow/compiler/mlir/stablehlo/BUILD index 355fd9ecb5d..0425d7d4300 100644 --- a/tensorflow/compiler/mlir/stablehlo/BUILD +++ b/tensorflow/compiler/mlir/stablehlo/BUILD @@ -32,6 +32,11 @@ tsl_pybind_extension( "-frtti", ], features = ["-use_header_modules"], + visibility = [ + ":friends", + "//tensorflow/python:__pkg__", + "//tensorflow/tools/pip_package:__subpackages__", + ], deps = [ "//third_party/python_runtime:headers", "@llvm-project//llvm:Support", diff --git a/tensorflow/compiler/mlir/tfr/BUILD b/tensorflow/compiler/mlir/tfr/BUILD index 333712a56e5..71cbd506612 100644 --- a/tensorflow/compiler/mlir/tfr/BUILD +++ b/tensorflow/compiler/mlir/tfr/BUILD @@ -272,6 +272,9 @@ tf_python_pybind_extension( pytype_srcs = [ "tfr_wrapper.pyi", ], + visibility = [ + "//tensorflow/python:__pkg__", + ], deps = [ ":tfr", "//tensorflow/compiler/mlir/tensorflow", diff --git a/tensorflow/compiler/mlir/tfr/build_defs.bzl b/tensorflow/compiler/mlir/tfr/build_defs.bzl index e9dd5e91780..9f10f82f0e1 100644 --- a/tensorflow/compiler/mlir/tfr/build_defs.bzl +++ b/tensorflow/compiler/mlir/tfr/build_defs.bzl @@ -1,9 +1,11 @@ """BUILD extension for TF composition project.""" +load("@local_tsl//third_party/py/rules_pywrap:pywrap.bzl", "use_pywrap_rules") load("//tensorflow:strict.default.bzl", "py_strict_binary", "py_strict_library") load("//tensorflow:tensorflow.bzl", "tf_custom_op_library", "tf_gen_op_wrapper_py") load("//tensorflow:tensorflow.default.bzl", "tf_custom_op_py_library") +# TODO(b/356020232): cleanup use_pywrap_rules once migration is done def gen_op_libraries( name, src, @@ -22,7 +24,11 @@ def gen_op_libraries( if not src.endswith(".py") or name == src[:-3]: fail("'src' %s conflicts with op Python wrapper. Rename it to be different from 'name'." % src) - py_deps = [ + py_deps = [] + if use_pywrap_rules(): + py_deps = ["//tensorflow/python:_pywrap_tensorflow"] + + py_deps += [ "//tensorflow/compiler/mlir/tfr:op_reg_gen", "//tensorflow/compiler/mlir/tfr:tfr_gen", "//tensorflow/compiler/mlir/tfr:composite", @@ -42,7 +48,9 @@ def gen_op_libraries( name = registered_op, srcs = [], outs = [name + ".inc.cc"], - cmd = "$(location %s) --output=$@ --gen_register_op=true" % gen_op_lib_exec, + cmd = + "PYWRAP_TARGET='//third_party/tensorflow/python:_pywrap_tensorflow' " + + "$(location %s) --output=$@ --gen_register_op=true" % gen_op_lib_exec, tools = [":" + gen_op_lib_exec], tags = tags, ) @@ -105,7 +113,9 @@ def gen_op_libraries( name = name + "_mlir", srcs = [], outs = [name + ".mlir"], - cmd = "$(location %s) --output=$@ --gen_register_op=false" % gen_tfr_lib_exec, + cmd = + "PYWRAP_TARGET='//third_party/tensorflow/python:_pywrap_tensorflow' " + + "$(location %s) --output=$@ --gen_register_op=false" % gen_tfr_lib_exec, tools = [":" + gen_tfr_lib_exec], tags = tags, ) diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index 8df04ec8e2b..3b7919ba3fe 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -1556,7 +1556,10 @@ cc_library( "tf2xla_opset.cc", ], hdrs = ["tf2xla_opset.h"], - visibility = ["//tensorflow/python:__pkg__"], + visibility = [ + "//tensorflow/python:__pkg__", + "//tensorflow/python/util:__pkg__", + ], deps = [ ":tf2xla_util", ":xla_op_registry", diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD index a729dddc9c5..455c19aa283 100644 --- a/tensorflow/core/BUILD +++ b/tensorflow/core/BUILD @@ -87,7 +87,16 @@ load( "tf_opts_nortti_if_lite_protos", "transitive_hdrs", ) -load("//tensorflow:tensorflow.default.bzl", "cc_header_only_library", "filegroup", "get_compatible_with_portable", "tensorflow_opensource_extra_deps", "tf_monitoring_framework_deps", "tf_selective_registration_deps") +load( + "//tensorflow:tensorflow.default.bzl", + "cc_header_only_library", + "custom_op_cc_header_only_library", + "filegroup", + "get_compatible_with_portable", + "tensorflow_opensource_extra_deps", + "tf_monitoring_framework_deps", + "tf_selective_registration_deps", +) # For platform specific build config load( @@ -1720,7 +1729,7 @@ tf_cuda_library( alwayslink = 1, ) -cc_header_only_library( +custom_op_cc_header_only_library( name = "framework_headers_lib", # Fully depend on external repositories, because identifying the headers # is fragile. diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD index 62bdeb4e005..304fe97d153 100644 --- a/tensorflow/core/framework/BUILD +++ b/tensorflow/core/framework/BUILD @@ -1570,6 +1570,7 @@ tf_proto_library( ":types_proto", ":versions_proto", ], + visibility = ["//visibility:public"], ) tf_proto_library( @@ -1802,6 +1803,9 @@ tf_proto_library( ":tensor_shape_proto", ":types_proto", ], + visibility = [ + "//tensorflow/python:__pkg__", + ] + default_visibility, ) # copybara:uncomment_begin(google-only) diff --git a/tensorflow/core/platform/build_config_root.bzl b/tensorflow/core/platform/build_config_root.bzl index 76fb425ba6f..744c9eada0d 100644 --- a/tensorflow/core/platform/build_config_root.bzl +++ b/tensorflow/core/platform/build_config_root.bzl @@ -9,6 +9,7 @@ load( _if_llvm_powerpc_available = "if_llvm_powerpc_available", _if_llvm_system_z_available = "if_llvm_system_z_available", _if_llvm_x86_available = "if_llvm_x86_available", + _if_pywrap = "if_pywrap", _if_static = "if_static", _if_static_and_not_mobile = "if_static_and_not_mobile", _tf_additional_grpc_deps_py = "tf_additional_grpc_deps_py", @@ -35,6 +36,7 @@ if_llvm_system_z_available = _if_llvm_system_z_available if_llvm_x86_available = _if_llvm_x86_available if_dynamic_kernels = _if_dynamic_kernels if_static = _if_static +if_pywrap = _if_pywrap if_static_and_not_mobile = _if_static_and_not_mobile tf_additional_grpc_deps_py = _tf_additional_grpc_deps_py tf_additional_license_deps = _tf_additional_license_deps diff --git a/tensorflow/core/platform/build_config_root.default.bzl b/tensorflow/core/platform/build_config_root.default.bzl index b503d99729b..bcea4eebe6d 100644 --- a/tensorflow/core/platform/build_config_root.default.bzl +++ b/tensorflow/core/platform/build_config_root.default.bzl @@ -1,5 +1,7 @@ """TODO(jakeharmon): Write module docstring.""" +load("@local_tsl//third_party/py/rules_pywrap:pywrap.bzl", "use_pywrap_rules") + # unused in TSL def tf_additional_plugin_deps(): return select({ @@ -10,6 +12,10 @@ def tf_additional_plugin_deps(): }) def if_dynamic_kernels(extra_deps, otherwise = []): + # TODO(b/356020232): remove after migration is done + if use_pywrap_rules(): + return otherwise + return select({ str(Label("//tensorflow:dynamic_loaded_kernels")): extra_deps, "//conditions:default": otherwise, diff --git a/tensorflow/dtensor/python/tests/BUILD b/tensorflow/dtensor/python/tests/BUILD index 80cb4361bcd..2803637bd85 100644 --- a/tensorflow/dtensor/python/tests/BUILD +++ b/tensorflow/dtensor/python/tests/BUILD @@ -394,6 +394,9 @@ dtensor_test( "tpu", ], main = "layout_test.py", + tags = [ + "no_windows", + ], deps = [ ":test_util", "//tensorflow/dtensor/python:api", diff --git a/tensorflow/lite/python/metrics/BUILD b/tensorflow/lite/python/metrics/BUILD index fdd6eb890fd..d213e57e3e9 100644 --- a/tensorflow/lite/python/metrics/BUILD +++ b/tensorflow/lite/python/metrics/BUILD @@ -16,7 +16,7 @@ cc_library( ), hdrs = ["wrapper/metrics_wrapper.h"], compatible_with = get_compatible_with_portable(), - visibility = ["//visibility:private"], + visibility = ["//tensorflow/python:__pkg__"], deps = [ "//third_party/python_runtime:headers", ] + if_portable( @@ -38,6 +38,7 @@ pybind_extension( ], visibility = [ "__subpackages__", + "//tensorflow/python:__pkg__", "//tensorflow/tools/pip_package:__subpackages__", ], deps = [ diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD index b42ff159414..0211c7abc83 100644 --- a/tensorflow/python/BUILD +++ b/tensorflow/python/BUILD @@ -11,7 +11,21 @@ load("//tensorflow:strict.default.bzl", "py_strict_library") # Placeholder: load py_proto_library load("//tensorflow:tensorflow.bzl", "VERSION", "cc_header_only_library", "clean_dep", "if_google", "if_oss", "if_windows", "if_xla_available", "tf_enable_mlir_bridge", "tf_python_pybind_static_deps", "tsl_async_value_deps") -load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "pywrap_tensorflow_macro", "tf_external_workspace_visible", "tf_monitoring_python_deps", "tf_pybind_cc_library_wrapper", "tf_python_pybind_extension") +load( + "//tensorflow:tensorflow.default.bzl", + "get_compatible_with_portable", + "pybind_extension", + "pywrap_aware_cc_import", + "pywrap_aware_filegroup", + "pywrap_aware_genrule", + "pywrap_common_library", + "pywrap_library", + "pywrap_tensorflow_macro", + "tf_external_workspace_visible", + "tf_monitoring_python_deps", + "tf_pybind_cc_library_wrapper", + "tf_python_pybind_extension", +) load( "//tensorflow/core/platform:build_config.bzl", "tf_additional_binary_deps", @@ -19,6 +33,7 @@ load( ) load( "//tensorflow/core/platform:build_config_root.bzl", + "if_pywrap", "if_static", "tf_additional_plugin_deps", "tf_additional_profiler_deps", @@ -873,7 +888,7 @@ pywrap_tensorflow_macro( # ** Targets for Windows build (start) ** # We need the following targets to expose symbols from _pywrap_tensorflow.dll -filegroup( +pywrap_aware_filegroup( name = "win_lib_files_for_exported_symbols", srcs = [ "//tensorflow/c:checkpoint_reader", # checkpoint_reader @@ -975,7 +990,7 @@ filegroup( # Filter the DEF file to reduce the number of symbols to 64K or less. # Note that we also write the name of the pyd file into DEF file so that # the dynamic libraries of custom ops can find it at runtime. -genrule( +pywrap_aware_genrule( name = "pywrap_tensorflow_filtered_def_file", srcs = select({ "//tensorflow:windows": [ @@ -1003,7 +1018,7 @@ genrule( ) # Write to a file a list of all cc_library targets that we need for exporting symbols on Windows. -genrule( +pywrap_aware_genrule( name = "pybind_symbol_target_libs_file", srcs = [":win_lib_files_for_exported_symbols"], outs = ["pybind_symbol_target_libs_file.txt"], @@ -1019,25 +1034,25 @@ genrule( ) # Get the import library of _pywrap_tensorflow_internal.pyd, platform-specific to Windows. -filegroup( +pywrap_aware_filegroup( name = "get_pywrap_tensorflow_import_lib_file", srcs = [":_pywrap_tensorflow_internal.so"], output_group = "interface_library", ) -cc_import( +pywrap_aware_cc_import( name = "_pywrap_tensorflow_internal_linux", shared_library = "//tensorflow/python:lib_pywrap_tensorflow_internal.so", visibility = tf_external_workspace_visible(visibility), ) -cc_import( +pywrap_aware_cc_import( name = "_pywrap_tensorflow_internal_macos", shared_library = "//tensorflow/python:lib_pywrap_tensorflow_internal.dylib", visibility = tf_external_workspace_visible(visibility), ) -cc_import( +pywrap_aware_cc_import( name = "_pywrap_tensorflow_internal_windows", interface_library = "//tensorflow/python:pywrap_tensorflow_import_lib_file", shared_library = "//tensorflow/python:_pywrap_tensorflow_internal.dll", @@ -1046,7 +1061,7 @@ cc_import( # Rename the import library for _pywrap_tensorflow_internal.pyd to _pywrap_tensorflow_internal.lib # (It was _pywrap_tensorflow_internal.so.if.lib). -genrule( +pywrap_aware_genrule( name = "pywrap_tensorflow_import_lib_file", srcs = [":get_pywrap_tensorflow_import_lib_file"], outs = ["_pywrap_tensorflow_internal.lib"], @@ -1059,7 +1074,7 @@ genrule( # Create a cc_import rule for the import library of _pywrap_tensorflow_internal.dll # so that custom ops' dynamic libraries can link against it. -cc_import( +pywrap_aware_cc_import( name = "pywrap_tensorflow_import_lib", interface_library = select({ "//tensorflow:windows": ":pywrap_tensorflow_import_lib_file", @@ -1109,7 +1124,7 @@ tf_python_pybind_extension( "//third_party/python_runtime:headers", "@com_google_absl//absl/strings", "@pybind11", - ], + ] + if_pywrap(["//tensorflow/compiler/mlir/python:mlir"]), ) py_strict_library( @@ -1198,24 +1213,26 @@ py_strict_library( tf_python_pybind_extension( name = "_pywrap_tfe", srcs = ["tfe_wrapper.cc"], - hdrs = [ - "//tensorflow/c:headers", - "//tensorflow/c:safe_ptr_hdr", - "//tensorflow/c/eager:headers", - "//tensorflow/c/eager:pywrap_required_hdrs", - "//tensorflow/c/experimental/ops:pywrap_required_hdrs", - "//tensorflow/core/common_runtime/eager:pywrap_required_hdrs", - "//tensorflow/core/distributed_runtime:pywrap_required_hdrs", - "//tensorflow/core/distributed_runtime/coordination:pywrap_required_hdrs", - "//tensorflow/core/distributed_runtime/eager:pywrap_required_hdrs", - "//tensorflow/python/eager:pywrap_required_hdrs", - "//tensorflow/python/lib/core:py_exception_registry_hdr", - "//tensorflow/python/lib/core:safe_pyobject_ptr_required_hdrs", - "//tensorflow/python/util:util_hdr", - "@local_xla//xla/tsl/distributed_runtime:pywrap_required_hdrs", - "@local_xla//xla/tsl/distributed_runtime/coordination:pywrap_required_hdrs", - "@local_xla//xla/tsl/python/lib/core:numpy_hdr", - ], + hdrs = if_pywrap( + if_false = [ + "//tensorflow/c:headers", + "//tensorflow/c:safe_ptr_hdr", + "//tensorflow/c/eager:headers", + "//tensorflow/c/eager:pywrap_required_hdrs", + "//tensorflow/c/experimental/ops:pywrap_required_hdrs", + "@local_xla//xla/tsl/distributed_runtime:pywrap_required_hdrs", + "@local_xla//xla/tsl/distributed_runtime/coordination:pywrap_required_hdrs", + "@local_xla//xla/tsl/python/lib/core:numpy_hdr", + "//tensorflow/core/common_runtime/eager:pywrap_required_hdrs", + "//tensorflow/core/distributed_runtime:pywrap_required_hdrs", + "//tensorflow/core/distributed_runtime/coordination:pywrap_required_hdrs", + "//tensorflow/core/distributed_runtime/eager:pywrap_required_hdrs", + "//tensorflow/python/eager:pywrap_required_hdrs", + "//tensorflow/python/lib/core:py_exception_registry_hdr", + "//tensorflow/python/lib/core:safe_pyobject_ptr_required_hdrs", + "//tensorflow/python/util:util_hdr", + ], + ), dynamic_deps = [":_pywrap_tensorflow_internal.so"] + select({ "//tensorflow:macos": ["//tensorflow:libtensorflow_framework.%s.dylib" % VERSION], "//conditions:default": ["//tensorflow:libtensorflow_framework.so.%s" % VERSION], @@ -1252,7 +1269,9 @@ tf_python_pybind_extension( "@com_google_absl//absl/types:optional", "@pybind11", # copybara:uncomment "@pybind11_protobuf//pybind11_protobuf:native_proto_caster", - ] + if_static( + ] + if_pywrap( + if_true = ["//tensorflow/python/eager:pywrap_tfe_lib"], + ) + if_static( extra_deps = [ "//tensorflow/core/protobuf:eager_service_proto_cc", "//tensorflow/core/protobuf:master_proto_cc", @@ -1429,3 +1448,175 @@ pytype_strict_library( "//tensorflow/python/util:tf_export", ], ) + +pywrap_library( + name = "_pywrap_tensorflow", + cc_deps_filter = [ + "@com_google_protobuf//:protobuf", + "@com_google_protobuf//:protobuf_lite", + "@zlib//:zlib", + ], + linkopts = select({ + "//tensorflow:windows": [ + "-DEFAULTLIB:ws2_32.lib", + "-DEFAULTLIB:advapi32.lib", + "-DEFAULTLIB:crypt32.lib", + "-DEFAULTLIB:Normaliz.lib", + ], + "//conditions:default": [], + }), + py_cc_deps_filter = select({ + "//tensorflow:windows": [], + "//conditions:default": [ + "@local_xla//xla/tsl/python/lib/core:ml_dtypes_lib", + "@local_xla//xla/tsl/python/lib/core:numpy", + "@local_xla//xla/backends/profiler/cpu:python_tracer_impl", + "@local_xla//xla/backends/profiler/cpu:python_tracer", + "@local_xla//xla/python/profiler/internal:python_hooks", + "//tensorflow/lite/python/interpreter_wrapper:python_error_reporter", + "//tensorflow/lite/python/interpreter_wrapper:python_utils", + "//tensorflow/lite/toco/python:toco_python_api", + "//tensorflow/python/client:tf_session_helper", + "//tensorflow/python/eager:pywrap_tfe_lib", + "//tensorflow/python/framework:op_def_util_cc", + "//tensorflow/python/framework:py_context_manager", + "//tensorflow/python/framework:python_api_info", + "//tensorflow/python/framework:python_api_parameter_converter", + "//tensorflow/python/framework:python_tensor_converter", + "//tensorflow/python/framework:python_api_dispatcher", + "//tensorflow/python/lib/core:ndarray_tensor_bridge", + "//tensorflow/python/lib/core:ndarray_tensor", + "//tensorflow/python/lib/core:py_seq_tensor", + "//tensorflow/python/lib/core:py_util", + "//tensorflow/python/lib/core:py_exception_registry", + "//tensorflow/python/lib/core:py_func_lib", + "//tensorflow/python/util:cpp_python_util", + "//tensorflow/python/util:function_parameter_canonicalizer", + "//tensorflow/python/util:stack_trace", + "//tensorflow/python/util:cpp_nest", + "//tensorflow/compiler/mlir/lite/python:converter_python_api", + "//tensorflow/lite/python/metrics:metrics_wrapper_lib", + "//tensorflow/lite/python/interpreter_wrapper:interpreter_wrapper_lib", + "//tensorflow/lite/python/interpreter_wrapper:numpy", + "//tensorflow/lite/python/optimize:calibration_wrapper_lib", + ], + }), + visibility = ["//visibility:public"], + # win_def_file = "_pywrap_tensorflow.def", + deps = [ + ":_pywrap_quantize_training", + ":_pywrap_tensorflow_cc_only", + "//tensorflow/compiler/mlir/lite/python:_pywrap_converter_api", + "//tensorflow/compiler/mlir/python/mlir_wrapper:filecheck_wrapper", + "//tensorflow/compiler/mlir/quantization/stablehlo/python:pywrap_quantization", + "//tensorflow/compiler/mlir/quantization/tensorflow/python:pywrap_function_lib", + "//tensorflow/compiler/mlir/quantization/tensorflow/python:pywrap_quantize_model", + "//tensorflow/compiler/mlir/stablehlo:stablehlo_extension", + "//tensorflow/compiler/mlir/tfr:tfr_wrapper", + "//tensorflow/compiler/tf2tensorrt:_pywrap_py_utils", + "//tensorflow/lite/python/analyzer_wrapper:_pywrap_analyzer_wrapper", + "//tensorflow/lite/python/interpreter_wrapper:_pywrap_tensorflow_interpreter_wrapper", + "//tensorflow/lite/python/metrics:_pywrap_tensorflow_lite_metrics_wrapper", + "//tensorflow/lite/python/optimize:_pywrap_tensorflow_lite_calibration_wrapper", + "//tensorflow/python:_pywrap_dtensor_device", + "//tensorflow/python:_pywrap_mlir", + "//tensorflow/python:_pywrap_parallel_device", + "//tensorflow/python:_pywrap_py_exception_registry", + "//tensorflow/python:_pywrap_sanitizers", + "//tensorflow/python:_pywrap_tfcompile", + "//tensorflow/python:_pywrap_tfe", + "//tensorflow/python:_pywrap_toco_api", + "//tensorflow/python:flags_pybind", + "//tensorflow/python/autograph/impl/testing:pybind_for_testing", + "//tensorflow/python/client:_pywrap_debug_events_writer", + "//tensorflow/python/client:_pywrap_device_lib", + "//tensorflow/python/client:_pywrap_events_writer", + "//tensorflow/python/client:_pywrap_tf_session", + "//tensorflow/python/data/experimental/service:_pywrap_server_lib", + "//tensorflow/python/data/experimental/service:_pywrap_snapshot_utils", + "//tensorflow/python/data/experimental/service:_pywrap_utils_exp", + "//tensorflow/python/framework:_dtypes", + "//tensorflow/python/framework:_errors_test_helper", + "//tensorflow/python/framework:_op_def_library_pybind", + "//tensorflow/python/framework:_op_def_registry", + "//tensorflow/python/framework:_op_def_util", + "//tensorflow/python/framework:_proto_comparators", + "//tensorflow/python/framework:_py_context_manager", + "//tensorflow/python/framework:_python_memory_checker_helper", + "//tensorflow/python/framework:_pywrap_python_api_dispatcher", + "//tensorflow/python/framework:_pywrap_python_api_info", + "//tensorflow/python/framework:_pywrap_python_api_parameter_converter", + "//tensorflow/python/framework:_pywrap_python_op_gen", + "//tensorflow/python/framework:_pywrap_python_tensor_converter", + "//tensorflow/python/framework:_test_metrics_util", + "//tensorflow/python/framework/experimental:_math_ops", + "//tensorflow/python/framework/experimental:_nn_ops", + "//tensorflow/python/framework/experimental:_tape", + "//tensorflow/python/framework/experimental:_unified_api", + "//tensorflow/python/grappler:_pywrap_cost_analyzer", + "//tensorflow/python/grappler:_pywrap_model_analyzer", + "//tensorflow/python/grappler:_pywrap_tf_cluster", + "//tensorflow/python/grappler:_pywrap_tf_item", + "//tensorflow/python/grappler:_pywrap_tf_optimizer", + "//tensorflow/python/lib/core:_pywrap_py_func", + "//tensorflow/python/lib/io:_pywrap_file_io", + "//tensorflow/python/lib/io:_pywrap_record_io", + "//tensorflow/python/platform:_pywrap_cpu_feature_guard", + "//tensorflow/python/platform:_pywrap_stacktrace_handler", + "//tensorflow/python/platform:_pywrap_tf2", + "//tensorflow/python/profiler/internal:_pywrap_profiler", + "//tensorflow/python/profiler/internal:_pywrap_traceme", + "//tensorflow/python/saved_model:pywrap_saved_model", + "//tensorflow/python/tpu:_pywrap_sparse_core_layout", + "//tensorflow/python/tpu:_pywrap_tpu_embedding", + "//tensorflow/python/util:_function_parameter_canonicalizer_binding_for_test", + "//tensorflow/python/util:_pywrap_checkpoint_reader", + "//tensorflow/python/util:_pywrap_determinism", + "//tensorflow/python/util:_pywrap_kernel_registry", + "//tensorflow/python/util:_pywrap_nest", + "//tensorflow/python/util:_pywrap_tensor_float_32_execution", + "//tensorflow/python/util:_pywrap_tfprof", + "//tensorflow/python/util:_pywrap_transform_graph", + "//tensorflow/python/util:_pywrap_util_port", + "//tensorflow/python/util:_pywrap_utils", + "//tensorflow/python/util:_tf_stack", + "//tensorflow/python/util:fast_module_type", + "//tensorflow/python/util:pywrap_xla_ops", + ], +) + +pybind_extension( + name = "_pywrap_tensorflow_cc_only", + srcs = [], + deps = [ + ":_protobuf_inline_symbols_enforcer", + "//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration", + "//tensorflow/core/distributed_runtime:server_lib", + "//tensorflow/core/distributed_runtime/rpc:grpc_server_lib", + "//tensorflow/core/distributed_runtime/rpc:grpc_session", + "//tensorflow/core/kernels:data_service_ops", + "//tensorflow/core/kernels:reader_ops", + "//tensorflow/distribute/experimental/rpc/kernels:rpc_ops", + "//tensorflow/dtensor/cc:tensor_layout", + "@local_xla//xla/backends/profiler/cpu:python_tracer", + ], +) + +cc_library( + name = "_protobuf_inline_symbols_enforcer", + srcs = ["protobuf_inline_symbols_enforcer.cc"], + deps = [ + "//tensorflow/compiler/mlir/quantization/tensorflow:exported_model_proto_cc", + "//tensorflow/core/framework:attr_value_proto_cc", + "//tensorflow/core/framework:function_proto_cc", + "//tensorflow/core/framework:graph_proto_cc", + "//tensorflow/core/protobuf:for_core_protos_cc", + "//tensorflow/dtensor/proto:layout_proto_cc", + "@local_tsl//tsl/profiler/protobuf:xplane_proto_cc", + ], +) + +pywrap_common_library( + name = "_pywrap_tensorflow_common", + dep = ":_pywrap_tensorflow", +) diff --git a/tensorflow/python/client/BUILD b/tensorflow/python/client/BUILD index ed9e8466ace..25055f4e24a 100644 --- a/tensorflow/python/client/BUILD +++ b/tensorflow/python/client/BUILD @@ -2,7 +2,7 @@ load("//tensorflow:strict.default.bzl", "py_strict_library") load("//tensorflow:tensorflow.bzl", "tf_cuda_library") load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test", "tf_python_pybind_extension") -load("//tensorflow/core/platform:build_config_root.bzl", "if_static") +load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap", "if_static") load( "//tensorflow/tools/test:performance.bzl", "cuda_py_benchmark_test", @@ -31,21 +31,29 @@ py_strict_library( tf_python_pybind_extension( name = "_pywrap_tf_session", srcs = ["tf_session_wrapper.cc"], - hdrs = [ - "tf_session_helper.h", - "//tensorflow/c:headers", - "//tensorflow/c:safe_ptr_hdr", - "//tensorflow/c/eager:headers", - "//tensorflow/c/eager:pywrap_required_hdrs", - "//tensorflow/c/experimental/ops:pywrap_required_hdrs", - "//tensorflow/core/common_runtime/eager:pywrap_required_hdrs", - "//tensorflow/core/distributed_runtime:pywrap_required_hdrs", - "//tensorflow/core/distributed_runtime/coordination:pywrap_required_hdrs", - "//tensorflow/core/distributed_runtime/eager:pywrap_required_hdrs", - "//tensorflow/python/lib/core:safe_pyobject_ptr_required_hdrs", - "@local_xla//xla/tsl/distributed_runtime:pywrap_required_hdrs", - "@local_xla//xla/tsl/distributed_runtime/coordination:pywrap_required_hdrs", - "@local_xla//xla/tsl/python/lib/core:numpy_hdr", + hdrs = if_pywrap( + if_false = [ + "tf_session_helper.h", + "//tensorflow/c:headers", + "//tensorflow/c:safe_ptr_hdr", + "//tensorflow/c/eager:headers", + "//tensorflow/c/eager:pywrap_required_hdrs", + "//tensorflow/c/experimental/ops:pywrap_required_hdrs", + "@local_xla//xla/tsl/distributed_runtime:pywrap_required_hdrs", + "@local_xla//xla/tsl/distributed_runtime/coordination:pywrap_required_hdrs", + "@local_xla//xla/tsl/python/lib/core:numpy_hdr", + "//tensorflow/core/common_runtime/eager:pywrap_required_hdrs", + "//tensorflow/core/distributed_runtime:pywrap_required_hdrs", + "//tensorflow/core/distributed_runtime/coordination:pywrap_required_hdrs", + "//tensorflow/core/distributed_runtime/eager:pywrap_required_hdrs", + "//tensorflow/python/lib/core:safe_pyobject_ptr_required_hdrs", + ], + if_true = [], + ), + additional_exported_symbols = [ + "_TF_SetTarget", + "_TF_SetConfig", + "_TF_NewSessionOptions", ], enable_stub_generation = True, pytype_srcs = [ @@ -72,7 +80,13 @@ tf_python_pybind_extension( "@pybind11", "@pybind11_abseil//pybind11_abseil:absl_casters", "@pybind11_protobuf//pybind11_protobuf:native_proto_caster", - ] + if_static( + ] + if_pywrap([ + "//tensorflow/c:safe_ptr", + "//tensorflow/c:c_api_experimental", + "//tensorflow/python/client:tf_session_helper", + "//tensorflow/c:python_api", + "//tensorflow/core/common_runtime:core_cpu_lib", + ]) + if_static( extra_deps = [ "//tensorflow/core/protobuf:eager_service_proto_cc", "//tensorflow/core/protobuf:master_proto_cc", diff --git a/tensorflow/python/framework/BUILD b/tensorflow/python/framework/BUILD index 5b3aef335e0..8df38adcfbf 100644 --- a/tensorflow/python/framework/BUILD +++ b/tensorflow/python/framework/BUILD @@ -16,9 +16,9 @@ load( "tf_gen_op_wrapper_py", "tf_kernel_library", ) -load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test", "tf_python_framework_friends", "tf_python_pybind_extension") +load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "stripped_cc_info", "tf_py_strict_test", "tf_python_framework_friends", "tf_python_pybind_extension") load("//tensorflow/core/platform:build_config.bzl", "pyx_library", "tf_additional_all_protos", "tf_additional_lib_deps", "tf_proto_library", "tf_protos_grappler") # @unused -load("//tensorflow/core/platform:build_config_root.bzl", "if_static", "tf_additional_xla_deps_py") +load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap", "if_static", "tf_additional_xla_deps_py") load("//tensorflow/python/tpu:tpu.bzl", "tpu_py_strict_test") load( "//tensorflow/tools/test:performance.bzl", @@ -33,8 +33,8 @@ package( licenses = ["notice"], ) -tf_cc_shared_object( - name = "test_file_system.so", +cc_library( + name = "test_file_system", srcs = ["test_file_system.cc"], copts = if_not_windows(["-Wno-sign-compare"]), linkopts = select({ @@ -45,12 +45,32 @@ tf_cc_shared_object( "//tensorflow:windows": [], }), deps = [ - "//tensorflow/core:framework_headers_lib", - "@com_google_protobuf//:protobuf_headers", - "@eigen_archive//:eigen3", + "//tensorflow/core:lib", + "//tensorflow/core/platform:null_file_system", ], ) +stripped_cc_info( + name = "test_file_system_stripped", + deps = [":test_file_system"], +) + +tf_cc_shared_object( + name = "test_file_system.so", + srcs = if_pywrap(if_false = ["test_file_system.cc"]), + copts = if_not_windows(["-Wno-sign-compare"]), + deps = if_pywrap( + if_false = [ + "@eigen_archive//:eigen3", + "//tensorflow/core:framework_headers_lib", + ], + if_true = [ + ":test_file_system_stripped", + "//tensorflow/python:_pywrap_tensorflow_common", + ], + ) + ["@com_google_protobuf//:protobuf_headers"], +) + tf_py_strict_test( name = "file_system_test", size = "small", @@ -656,8 +676,10 @@ tf_python_pybind_extension( pytype_srcs = [ "_op_def_library_pybind.pyi", ], - deps = [ - ":op_def_util_headers", + deps = if_pywrap( + if_false = [":op_def_util_headers"], + if_true = [":op_def_util_cc"], + ) + [ "//tensorflow/core:framework_headers_lib", "//tensorflow/core:lib_headers_for_pybind", "//tensorflow/core:lib_proto_parsing", @@ -796,8 +818,10 @@ tf_python_pybind_extension( pytype_srcs = [ "_op_def_util.pyi", ], - deps = [ - ":op_def_util_headers", + deps = if_pywrap( + if_false = [":op_def_util_headers"], + if_true = [":op_def_util_cc"], + ) + [ "//tensorflow/core:lib_headers_for_pybind", "//tensorflow/core:protos_all_cc", "//tensorflow/python/lib/core:pybind11_lib", @@ -852,10 +876,6 @@ tf_python_pybind_extension( name = "_pywrap_python_api_parameter_converter", srcs = ["python_api_parameter_converter_wrapper.cc"], hdrs = [ - "op_def_util.h", - "python_api_info.h", - "python_api_parameter_converter.h", - "python_tensor_converter.h", "//tensorflow/c:headers", "//tensorflow/c/eager:pywrap_required_hdrs", "//tensorflow/c/experimental/ops:pywrap_required_hdrs", @@ -868,7 +888,14 @@ tf_python_pybind_extension( "@local_xla//xla/tsl/distributed_runtime:pywrap_required_hdrs", "@local_xla//xla/tsl/distributed_runtime/coordination:pywrap_required_hdrs", "@local_xla//xla/tsl/python/lib/core:numpy_hdr", - ], + ] + if_pywrap( + if_false = [ + "op_def_util.h", + "python_api_info.h", + "python_api_parameter_converter.h", + "python_tensor_converter.h", + ], + ), enable_stub_generation = True, pytype_srcs = [ "_pywrap_python_api_parameter_converter.pyi", @@ -890,7 +917,14 @@ tf_python_pybind_extension( "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", "@pybind11", - ] + if_static( + ] + if_pywrap( + if_true = [ + ":python_api_parameter_converter", + ":python_api_info", + ":op_def_util_cc", + ":python_tensor_converter", + ], + ) + if_static( extra_deps = [ "//tensorflow/core/protobuf:eager_service_proto_cc", "//tensorflow/core/protobuf:master_proto_cc", @@ -953,9 +987,6 @@ tf_python_pybind_extension( name = "_pywrap_python_api_info", srcs = ["python_api_info_wrapper.cc"], hdrs = [ - "op_def_util.h", - "python_api_info.h", - "python_tensor_converter.h", "//tensorflow/c:headers", "//tensorflow/c/eager:pywrap_required_hdrs", "//tensorflow/c/experimental/ops:pywrap_required_hdrs", @@ -968,7 +999,13 @@ tf_python_pybind_extension( "@local_xla//xla/tsl/distributed_runtime:pywrap_required_hdrs", "@local_xla//xla/tsl/distributed_runtime/coordination:pywrap_required_hdrs", "@local_xla//xla/tsl/python/lib/core:numpy_hdr", - ], + ] + if_pywrap( + if_false = [ + "op_def_util.h", + "python_api_info.h", + "python_tensor_converter.h", + ], + ), enable_stub_generation = True, pytype_srcs = [ "_pywrap_python_api_info.pyi", @@ -990,7 +1027,13 @@ tf_python_pybind_extension( "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", "@pybind11", - ] + if_static( + ] + if_pywrap( + if_true = [ + ":python_api_info", + ":op_def_util_cc", + ":python_tensor_converter", + ], + ) + if_static( extra_deps = [ "@local_xla//xla/tsl/protobuf:coordination_service_proto_cc", "//tensorflow/core/protobuf:eager_service_proto_cc", @@ -1064,7 +1107,7 @@ tf_python_pybind_extension( "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", "@pybind11", - ], + ] + if_pywrap([":python_api_dispatcher"]), ) tf_py_strict_test( diff --git a/tensorflow/python/framework/experimental/BUILD b/tensorflow/python/framework/experimental/BUILD index 2d7a8f11129..273cf42c4e1 100644 --- a/tensorflow/python/framework/experimental/BUILD +++ b/tensorflow/python/framework/experimental/BUILD @@ -2,6 +2,7 @@ load("//tensorflow:strict.default.bzl", "py_strict_library") load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_python_pybind_extension") +load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap") load( "//tensorflow/tools/test:performance.bzl", "cuda_py_benchmark_test", @@ -20,6 +21,9 @@ tf_python_pybind_extension( pytype_srcs = [ "_unified_api.pyi", ], + visibility = [ + "//tensorflow/python:__pkg__", + ], deps = [ "//tensorflow/c/eager:tfe_tensorhandle_internal", "//tensorflow/core:lib", @@ -35,6 +39,9 @@ tf_python_pybind_extension( name = "_tape", srcs = ["tape.cc"], features = ["-layering_check"], + visibility = [ + "//tensorflow/python:__pkg__", + ], deps = [ "//tensorflow/c/eager:tfe_tensorhandle_internal", "//tensorflow/core:lib", @@ -43,7 +50,13 @@ tf_python_pybind_extension( "//tensorflow/python:unified_api_pywrap_required_headers", "//tensorflow/python/lib/core:pybind11_lib", "@pybind11", - ], + ] + if_pywrap( + if_true = [ + "//tensorflow/c/experimental/gradients/tape:tape_context", + "//tensorflow/c/experimental/gradients:math_grad", + "//tensorflow/c/experimental/gradients:nn_grad", + ], + ), ) tf_python_pybind_extension( @@ -53,6 +66,9 @@ tf_python_pybind_extension( pytype_srcs = [ "_math_ops.pyi", ], + visibility = [ + "//tensorflow/python:__pkg__", + ], deps = [ "//tensorflow/c/eager:tfe_tensorhandle_internal", "//tensorflow/core:framework", @@ -63,7 +79,11 @@ tf_python_pybind_extension( "//tensorflow/python/lib/core:pybind11_lib", "@com_google_absl//absl/types:span", "@pybind11", - ], + ] + if_pywrap( + if_true = [ + "//tensorflow/c/experimental/ops:math_ops", + ], + ), ) tf_python_pybind_extension( @@ -73,6 +93,9 @@ tf_python_pybind_extension( pytype_srcs = [ "_nn_ops.pyi", ], + visibility = [ + "//tensorflow/python:__pkg__", + ], deps = [ "//tensorflow/c/eager:tfe_tensorhandle_internal", "//tensorflow/core:framework", @@ -83,7 +106,11 @@ tf_python_pybind_extension( "//tensorflow/python/lib/core:pybind11_lib", "@com_google_absl//absl/types:span", "@pybind11", - ], + ] + if_pywrap( + if_true = [ + "//tensorflow/c/experimental/ops:nn_ops", + ], + ), ) py_strict_library( diff --git a/tensorflow/python/grappler/BUILD b/tensorflow/python/grappler/BUILD index 366ebfa1927..1e1d643602b 100644 --- a/tensorflow/python/grappler/BUILD +++ b/tensorflow/python/grappler/BUILD @@ -2,6 +2,7 @@ load("//tensorflow:strict.default.bzl", "py_strict_binary", "py_strict_library") load("//tensorflow:tensorflow.bzl", "if_not_windows") load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "get_compatible_with_portable", "tf_py_strict_test", "tf_pybind_cc_library_wrapper", "tf_python_pybind_extension") load("//tensorflow/core/platform:build_config.bzl", "tf_protos_grappler") +load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap") package( # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], @@ -96,7 +97,7 @@ tf_python_pybind_extension( "//tensorflow/core:protos_all_cc", "//tensorflow/python/lib/core:pybind11_status", "@pybind11", - ], + ] + if_pywrap(["//tensorflow/python/grappler:model_analyzer_lib"]), ) py_strict_library( @@ -201,11 +202,14 @@ tf_python_pybind_extension( srcs = ["cluster_wrapper.cc"], hdrs = [ "//tensorflow/cc:pywrap_required_hdrs", - "//tensorflow/core/grappler:pywrap_required_hdrs", - "//tensorflow/core/grappler/clusters:pywrap_required_hdrs", - "//tensorflow/core/grappler/costs:pywrap_required_hdrs", - "//tensorflow/core/grappler/utils:pywrap_required_hdrs", - ], + ] + if_pywrap( + if_false = [ + "//tensorflow/core/grappler:pywrap_required_hdrs", + "//tensorflow/core/grappler/clusters:pywrap_required_hdrs", + "//tensorflow/core/grappler/costs:pywrap_required_hdrs", + "//tensorflow/core/grappler/utils:pywrap_required_hdrs", + ], + ), enable_stub_generation = True, pytype_srcs = [ "_pywrap_tf_cluster.pyi", @@ -219,7 +223,12 @@ tf_python_pybind_extension( "//tensorflow/python/lib/core:pybind11_status", "@com_google_absl//absl/types:span", "@pybind11", - ], + ] + if_pywrap( + if_true = [ + "//tensorflow/core/grappler/costs:measuring_cost_estimator", + "//tensorflow/core/grappler/clusters:single_machine", + ], + ), ) cuda_py_strict_test( @@ -265,14 +274,16 @@ py_strict_library( tf_python_pybind_extension( name = "_pywrap_tf_optimizer", srcs = ["tf_optimizer_wrapper.cc"], - hdrs = [ - "//tensorflow/cc:pywrap_required_hdrs", - "//tensorflow/core/grappler:pywrap_required_hdrs", - "//tensorflow/core/grappler/clusters:pywrap_required_hdrs", - "//tensorflow/core/grappler/costs:pywrap_required_hdrs", - "//tensorflow/core/grappler/optimizers:pywrap_required_hdrs", - "//tensorflow/core/grappler/verifiers:pywrap_required_hdrs", - ], + hdrs = if_pywrap( + if_false = [ + "//tensorflow/cc:pywrap_required_hdrs", + "//tensorflow/core/grappler:pywrap_required_hdrs", + "//tensorflow/core/grappler/clusters:pywrap_required_hdrs", + "//tensorflow/core/grappler/costs:pywrap_required_hdrs", + "//tensorflow/core/grappler/optimizers:pywrap_required_hdrs", + "//tensorflow/core/grappler/verifiers:pywrap_required_hdrs", + ], + ), enable_stub_generation = True, pytype_srcs = [ "_pywrap_tf_optimizer.pyi", @@ -293,7 +304,16 @@ tf_python_pybind_extension( "//tensorflow/python/lib/core:pybind11_status", "@pybind11", "@pybind11_protobuf//pybind11_protobuf:native_proto_caster", - ], + ] + if_pywrap( + if_true = [ + "//tensorflow/core/grappler/clusters:cluster", + "//tensorflow/core/grappler/clusters:utils", + "//tensorflow/core/grappler:grappler_item_builder", + "//tensorflow/core/grappler/optimizers:meta_optimizer", + "//tensorflow/core/grappler/optimizers:graph_optimizer", + "//tensorflow/core/grappler/verifiers:graph_verifier", + ], + ), ) tf_py_strict_test( diff --git a/tensorflow/python/kernel_tests/proto/BUILD b/tensorflow/python/kernel_tests/proto/BUILD index e50955770c1..12979717321 100644 --- a/tensorflow/python/kernel_tests/proto/BUILD +++ b/tensorflow/python/kernel_tests/proto/BUILD @@ -3,10 +3,10 @@ load("//tensorflow:strict.default.bzl", "py_strict_library") # Placeholder: load py_proto_library -load("//tensorflow:tensorflow.bzl", "tf_cc_shared_object") -load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test") +load("//tensorflow:tensorflow.bzl", "if_oss", "tf_cc_shared_object") +load("//tensorflow:tensorflow.default.bzl", "stripped_cc_info", "tf_py_strict_test") load("//tensorflow/core/platform:build_config.bzl", "tf_additional_all_protos", "tf_proto_library") -load("//tensorflow/core/platform:build_config_root.bzl", "if_static") +load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap") package( # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], @@ -22,10 +22,7 @@ tf_py_strict_test( name = "decode_proto_op_test", size = "small", srcs = ["decode_proto_op_test.py"], - data = if_static( - [], - otherwise = [":libtestexample.so"], - ), + data = if_oss([":libtestexample.so"]), python_version = "PY3", tags = [ "no_pip", # TODO(b/78026780) @@ -43,10 +40,7 @@ tf_py_strict_test( name = "encode_proto_op_test", size = "small", srcs = ["encode_proto_op_test.py"], - data = if_static( - [], - otherwise = [":libtestexample.so"], - ), + data = if_oss([":libtestexample.so"]), python_version = "PY3", tags = [ "no_pip", # TODO(b/78026780) @@ -118,9 +112,20 @@ tf_proto_library( tf_cc_shared_object( name = "libtestexample.so", linkstatic = 1, - deps = [ - ":test_example_proto_cc", - ], + deps = if_pywrap( + if_false = [ + ":test_example_proto_cc", + ], + if_true = [ + "//tensorflow/python:_pywrap_tensorflow_common", + ":test_example_proto_cc_stripped", + ], + ), +) + +stripped_cc_info( + name = "test_example_proto_cc_stripped", + deps = [":test_example_proto_cc"], ) py_strict_library( diff --git a/tensorflow/python/profiler/internal/BUILD b/tensorflow/python/profiler/internal/BUILD index 7fe4593d15c..a64bfd314d5 100644 --- a/tensorflow/python/profiler/internal/BUILD +++ b/tensorflow/python/profiler/internal/BUILD @@ -107,6 +107,7 @@ tf_python_pybind_extension( ], visibility = [ "//perftools/accelerators/xprof/xprofilez/integration_tests:__pkg__", + "//tensorflow/python:__pkg__", "//tensorflow/python/profiler:__subpackages__", "//tensorflow/tools/pip_package:__subpackages__", ], @@ -126,6 +127,7 @@ tf_python_pybind_extension( ], visibility = [ "//tensorflow/core/profiler:internal", + "//tensorflow/python:__pkg__", "//tensorflow/python/eager:__pkg__", "//tensorflow/python/profiler:__pkg__", "//tensorflow/tools/pip_package:__subpackages__", diff --git a/tensorflow/python/protobuf_inline_symbols_enforcer.cc b/tensorflow/python/protobuf_inline_symbols_enforcer.cc new file mode 100644 index 00000000000..24beeeb70fd --- /dev/null +++ b/tensorflow/python/protobuf_inline_symbols_enforcer.cc @@ -0,0 +1,91 @@ +/* Copyright 2021 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#include + +#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h" +#include "tensorflow/core/framework/attr_value.pb.h" +#include "tensorflow/core/framework/function.pb.h" +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/protobuf/config.pb.h" +#include "tensorflow/core/protobuf/data_service.pb.h" +#include "tensorflow/core/protobuf/device_properties.pb.h" +#include "tensorflow/core/protobuf/meta_graph.pb.h" +#include "tensorflow/core/protobuf/service_config.pb.h" +#include "tensorflow/dtensor/proto/layout.pb.h" +#include "tsl/profiler/protobuf/xplane.pb.h" + +namespace tensorflow { +namespace python { +void protobuf_inline_symbols_enforcer() { + tensorflow::NamedDevice named_device; + named_device.mutable_properties(); + named_device.properties(); + + tensorflow::NamedDevice named_device_move(std::move(named_device)); + named_device_move.mutable_properties(); + + tensorflow::quantization::ExportedModel exported_model; + exported_model.function_aliases(); + + tensorflow::profiler::XSpace x_space; + x_space.mutable_hostnames(); + x_space.mutable_hostnames(0); + + tensorflow::dtensor::LayoutProto layout_proto; + layout_proto.GetDescriptor(); + layout_proto.GetReflection(); + layout_proto.default_instance(); + + tensorflow::dtensor::MeshProto mesh_proto; + mesh_proto.GetDescriptor(); + mesh_proto.GetReflection(); + mesh_proto.default_instance(); + + tensorflow::FunctionDef function_def; + function_def.descriptor(); + function_def.GetDescriptor(); + function_def.GetReflection(); + function_def.default_instance(); + + tensorflow::FunctionDefLibrary function_def_library; + function_def_library.descriptor(); + + tensorflow::GraphDef graph_def; + graph_def.descriptor(); + graph_def.GetDescriptor(); + graph_def.GetReflection(); + graph_def.default_instance(); + + tensorflow::MetaGraphDef meta_graph_def; + meta_graph_def.GetDescriptor(); + meta_graph_def.GetReflection(); + meta_graph_def.default_instance(); + + tensorflow::AttrValue attr_value; + attr_value.default_instance(); + + tensorflow::ConfigProto config_proto; + config_proto.default_instance(); + + tensorflow::data::experimental::DispatcherConfig dispatcher_config; + dispatcher_config.default_instance(); + + tensorflow::data::experimental::WorkerConfig worker_config; + worker_config.default_instance(); + + tensorflow::data::DataServiceMetadata data_service_metadata; +} +} // namespace python +} // namespace tensorflow diff --git a/tensorflow/python/saved_model/BUILD b/tensorflow/python/saved_model/BUILD index 12ad8fadc6c..acd8241c9c6 100644 --- a/tensorflow/python/saved_model/BUILD +++ b/tensorflow/python/saved_model/BUILD @@ -927,6 +927,7 @@ tf_python_pybind_extension( "pywrap_saved_model/metrics.pyi", ], visibility = [ + "//tensorflow/python:__pkg__", "//tensorflow/python/checkpoint:__subpackages__", "//tensorflow/python/tpu:__pkg__", "//tensorflow/python/training:__subpackages__", diff --git a/tensorflow/python/tools/BUILD b/tensorflow/python/tools/BUILD index 324fcb83897..b9f0628f638 100644 --- a/tensorflow/python/tools/BUILD +++ b/tensorflow/python/tools/BUILD @@ -3,6 +3,7 @@ load("//tensorflow:strict.default.bzl", "py_strict_binary", "py_strict_library", "py_strict_test") load("//tensorflow:tensorflow.bzl", "if_google", "if_xla_available", "tf_cc_test") +load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap") load("//tensorflow/python/tools:tools.bzl", "saved_model_compile_aot") package( @@ -97,7 +98,9 @@ py_strict_binary( srcs = ["freeze_graph.py"], python_version = "PY3", srcs_version = "PY3", - deps = [":freeze_graph_lib"], + deps = [":freeze_graph_lib"] + if_pywrap( + if_true = ["//tensorflow/python:_pywrap_tensorflow"], + ), ) py_strict_binary( @@ -356,7 +359,9 @@ py_strict_binary( srcs = ["saved_model_cli.py"], python_version = "PY3", srcs_version = "PY3", - deps = [":saved_model_cli_lib"], + deps = [":saved_model_cli_lib"] + if_pywrap( + if_true = ["//tensorflow/python:_pywrap_tensorflow"], + ), ) py_strict_library( @@ -471,7 +476,7 @@ py_strict_binary( "//tensorflow/python/trackable:autotrackable", "@absl_py//absl:app", "@absl_py//absl/flags", - ], + ] + if_pywrap(["//tensorflow/python:_pywrap_tensorflow"]), ) # copybara:comment_begin(oss-only) @@ -497,6 +502,7 @@ genrule( name = "create_models_for_aot_compile", outs = EMITTED_AOT_SAVE_MODEL_OBJECTS, cmd = ( + "PYWRAP_TARGET='//third_party/tensorflow/python:_pywrap_tensorflow' " + "$(location :make_aot_compile_models) --out_dir $(@D)" ), tags = ["no_rocm"], diff --git a/tensorflow/python/tools/tools.bzl b/tensorflow/python/tools/tools.bzl index 0ed7102674b..42a95ef19da 100644 --- a/tensorflow/python/tools/tools.bzl +++ b/tensorflow/python/tools/tools.bzl @@ -132,6 +132,7 @@ def saved_model_compile_aot( "{}_makefile.inc".format(name), ], cmd = ( + "PYWRAP_TARGET='//third_party/tensorflow/python:_pywrap_tensorflow' " + "$(location {}) aot_compile_cpu ".format( clean_dep("//tensorflow/python/tools:saved_model_cli"), ) + diff --git a/tensorflow/python/tpu/BUILD b/tensorflow/python/tpu/BUILD index d7c88da50a9..4c857a6b724 100644 --- a/tensorflow/python/tpu/BUILD +++ b/tensorflow/python/tpu/BUILD @@ -7,6 +7,7 @@ load("//tensorflow:strict.default.bzl", "py_strict_library", "py_strict_test") # Placeholder: load py_proto_library load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test", "tf_python_pybind_extension") load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library") +load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap") load("//tensorflow/python/tpu:tpu.bzl", "internal_create_sanitizer_settings", "tpu_py_strict_test") # Do not add anymore paths here. You do not need to be in the visibility list @@ -1027,7 +1028,6 @@ tf_python_pybind_extension( "_pywrap_sparse_core_layout.pyi", ], deps = [ - "//tensorflow/core/tpu/kernels:_pywrap_sparse_core_layout_header_only", "//tensorflow/python/lib/core:pybind11_lib", "//tensorflow/python/lib/core:pybind11_status", "//tensorflow/python/lib/core:pybind11_status_headers", @@ -1035,7 +1035,14 @@ tf_python_pybind_extension( "@pybind11", "@pybind11_abseil//pybind11_abseil:status_casters", "@pybind11_protobuf//pybind11_protobuf:native_proto_caster", - ], + ] + if_pywrap( + if_false = [ + "//tensorflow/core/tpu/kernels:_pywrap_sparse_core_layout_header_only", + ], + if_true = [ + "//tensorflow/core/tpu/kernels:sparse_core_layout", + ], + ), ) tf_python_pybind_extension( diff --git a/tensorflow/python/util/BUILD b/tensorflow/python/util/BUILD index ee56bb821a2..9933709da15 100644 --- a/tensorflow/python/util/BUILD +++ b/tensorflow/python/util/BUILD @@ -4,7 +4,7 @@ load("//tensorflow:pytype.default.bzl", "pytype_strict_library") load("//tensorflow:strict.default.bzl", "py_strict_library") load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "tf_py_strict_test", "tf_python_pybind_extension") load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library") # @unused -load("//tensorflow/core/platform:build_config_root.bzl", "if_static") +load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap", "if_static") visibility = [ "//engedu/ml/tf_from_scratch:__pkg__", @@ -72,10 +72,22 @@ tf_python_pybind_extension( pytype_srcs = [ "_pywrap_tfprof.pyi", ], - deps = [ - "//tensorflow/core/profiler/internal:print_model_analysis_hdr", - "@pybind11", - ], + deps = if_pywrap( + if_false = [ + "//tensorflow/core/profiler/internal:print_model_analysis_hdr", + "@pybind11", + ], + if_true = [ + "//tensorflow/core:framework", + "//tensorflow/core/framework:reader_base", + "//tensorflow/core:lib_headers_for_pybind", + "//tensorflow/core/profiler/internal:print_model_analysis", + "//third_party/python_runtime:headers", + "@com_google_absl//absl/strings", + "@eigen_archive//:eigen3", + "@pybind11", + ], + ), ) tf_python_pybind_extension( @@ -107,7 +119,7 @@ tf_python_pybind_extension( "//tensorflow/python/lib/core:pybind11_lib", "//third_party/python_runtime:headers", "@pybind11", - ], + ] + if_pywrap([":cpp_nest"]), ) cc_library( @@ -140,7 +152,7 @@ tf_python_pybind_extension( "//tensorflow/python/lib/core:pybind11_lib", "//third_party/python_runtime:headers", "@pybind11", - ], + ] + if_pywrap([":kernel_registry"]), ) tf_python_pybind_extension( @@ -257,7 +269,7 @@ tf_python_pybind_extension( "//third_party/py/numpy:headers", "//third_party/python_runtime:headers", "@pybind11", - ], + ] + if_pywrap(["@com_google_absl//absl/strings"]), ) cc_library( @@ -527,7 +539,7 @@ cc_library( tf_python_pybind_extension( name = "_function_parameter_canonicalizer_binding_for_test", - testonly = True, + # testonly = True, srcs = ["function_parameter_canonicalizer_binding_for_test.cc"], hdrs = [ "function_parameter_canonicalizer.h", @@ -543,7 +555,12 @@ tf_python_pybind_extension( "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/types:span", "@pybind11", - ], + ] + if_pywrap( + if_true = [ + "//tensorflow/compiler/tf2xla:tf2xla_opset", + "//tensorflow/python/lib/core:pybind11_lib", + ], + ), ) tf_py_strict_test( diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl index 7c15609ae71..86c8e7df662 100644 --- a/tensorflow/tensorflow.bzl +++ b/tensorflow/tensorflow.bzl @@ -69,6 +69,7 @@ load( "@local_xla//xla/tsl:tsl.bzl", "tsl_gpu_library", _cc_header_only_library = "cc_header_only_library", + _custom_op_cc_header_only_library = "custom_op_cc_header_only_library", _if_cuda_or_rocm = "if_cuda_or_rocm", _if_cuda_tools = "if_cuda_tools", _if_nccl = "if_nccl", @@ -79,6 +80,12 @@ load( "if_tensorrt", "if_tensorrt_exec", ) +load( + "@local_tsl//third_party/py/rules_pywrap:pywrap.bzl", + "use_pywrap_rules", + _pybind_extension = "pybind_extension", + _stripped_cc_info = "stripped_cc_info", +) # Do not sort: copybara rule changes this def register_extension_info(**kwargs): @@ -109,6 +116,7 @@ def clean_dep(target): return str(Label(target)) cc_header_only_library = _cc_header_only_library +custom_op_cc_header_only_library = _custom_op_cc_header_only_library transitive_hdrs = _transitive_hdrs def if_oss(oss_value, google_value = []): @@ -656,9 +664,13 @@ def _rpath_user_link_flags(name): ], }) +# TODO(b/356020232): remove completely after migration is done # Bazel-generated shared objects which must be linked into TensorFlow binaries # to define symbols from //tensorflow/core:framework and //tensorflow/core:lib. def tf_binary_additional_srcs(fullversion = False): + if use_pywrap_rules(): + return [] + if fullversion: suffix = "." + VERSION else: @@ -674,7 +686,11 @@ def tf_binary_additional_srcs(fullversion = False): ], ) +# TODO(b/356020232): remove completely after migration is done def tf_binary_additional_data_deps(): + if use_pywrap_rules(): + return [] + return if_static( extra_deps = [], macos = [ @@ -689,7 +705,11 @@ def tf_binary_additional_data_deps(): ], ) +# TODO(b/356020232): remove completely after migration is done def tf_binary_pybind_deps(): + if use_pywrap_rules(): + return [] + return select({ clean_dep("//tensorflow:macos"): [ clean_dep( @@ -708,8 +728,12 @@ def tf_binary_pybind_deps(): ], }) +# TODO(b/356020232): remove completely after migration is done # Helper function for the per-OS tensorflow libraries and their version symlinks def tf_shared_library_deps(): + if use_pywrap_rules(): + return [] + return select({ clean_dep("//tensorflow:macos_with_framework_shared_object"): [ clean_dep("//tensorflow:libtensorflow.dylib"), @@ -775,6 +799,11 @@ def tf_cc_shared_object( visibility = None, **kwargs): """Configure the shared object (.so) file for TensorFlow.""" + + actual_framework_so = framework_so + if use_pywrap_rules(): + actual_framework_so = [] + if soversion != None: suffix = "." + str(soversion).split(".")[0] longsuffix = "." + str(soversion) @@ -825,13 +854,13 @@ def tf_cc_shared_object( soname = name_os_major.split("/")[-1] data_extra = [] - if framework_so != []: + if actual_framework_so != []: data_extra = tf_binary_additional_data_deps() cc_binary( exec_properties = if_google({"cpp_link.mem": "16g"}, {}), name = name_os_full, - srcs = srcs + framework_so, + srcs = srcs + actual_framework_so, deps = deps, linkshared = 1, data = data + data_extra, @@ -865,6 +894,7 @@ def tf_cc_shared_object( testonly = testonly, ) +# TODO(b/356020232): remove completely after migration is done # buildozer: disable=function-docstring-args def tf_cc_shared_library_opensource( name, @@ -885,6 +915,10 @@ def tf_cc_shared_library_opensource( win_def_file = None, visibility = None): """Configures the shared object file for TensorFlow.""" + + if use_pywrap_rules(): + return + names = _get_shared_library_name_os_version_matrix( name, per_os_targets = per_os_targets, @@ -944,6 +978,7 @@ def tf_cc_shared_library_opensource( visibility = visibility, ) +# TODO(b/356020232): remove completely after migration is done def _tf_cc_shared_library_opensource( name, additional_linker_inputs = None, @@ -960,6 +995,9 @@ def _tf_cc_shared_library_opensource( user_link_flags = None, visibility = None, win_def_file = None): + if use_pywrap_rules(): + return + cc_library_name = name + "_cclib" cc_library( name = cc_library_name, @@ -2058,13 +2096,17 @@ def tf_kernel_library( ) # TODO(gunan): CUDA dependency not clear here. Fix it. - tf_cc_shared_object( - name = "libtfkernel_%s.so" % name, - srcs = srcs + hdrs + textual_hdrs, - copts = copts, - tags = ["manual", "notap"], - deps = deps, - ) + # TODO(b/356020232): remove completely after migration is done + if use_pywrap_rules(): + pass + else: + tf_cc_shared_object( + name = "libtfkernel_%s.so" % name, + srcs = srcs + hdrs + textual_hdrs, + copts = copts, + tags = ["manual", "notap"], + deps = deps, + ) register_extension_info( extension = tf_kernel_library, @@ -2256,6 +2298,7 @@ check_deps = rule( }, ) +# TODO(b/356020232): cleanup use_pywrap_rules after migration is done def tf_custom_op_library( name, srcs = [], @@ -2270,14 +2313,26 @@ def tf_custom_op_library( if not gpu_deps: gpu_deps = [] - deps = deps + if_cuda_or_rocm([ + if use_pywrap_rules(): + deps = [clean_dep("//tensorflow/python:_pywrap_tensorflow_common")] + deps + else: + deps = list(deps) + + deps += if_cuda_or_rocm([ clean_dep("//tensorflow/core:stream_executor_headers_lib"), ]) + if_cuda([ "@local_config_cuda//cuda:cuda_headers", "@local_config_cuda//cuda:cuda_runtime", - ]) + if_windows([ - clean_dep("//tensorflow/python:pywrap_tensorflow_import_lib"), - ]) + tf_custom_op_library_additional_deps() + ]) + + if use_pywrap_rules(): + pass + else: + deps += if_windows( + [clean_dep("//tensorflow/python:pywrap_tensorflow_import_lib")], + ) + + deps += tf_custom_op_library_additional_deps() # Override EIGEN_STRONG_INLINE to inline when # --define=override_eigen_strong_inline=true to avoid long compiling time. @@ -2386,6 +2441,7 @@ _append_init_to_versionscript = rule( implementation = _append_init_to_versionscript_impl, ) +# TODO(b/356020232): remove completely after migration is done # This macro should only be used for pywrap_tensorflow_internal.so. # It was copied and refined from the original tf_py_wrap_cc_opensource rule. # buildozer: disable=function-docstring-args @@ -2401,6 +2457,15 @@ def pywrap_tensorflow_macro_opensource( version_script = None, win_def_file = None): """Builds the pywrap_tensorflow_internal shared object.""" + + if use_pywrap_rules(): + native.py_library( + name = name, + srcs = [], + deps = [], + ) + return + module_name = name.split("/")[-1] # Convert a rule name such as foo/bar/baz to foo/bar/_baz.so @@ -2531,6 +2596,8 @@ def pywrap_tensorflow_macro_opensource( # Export open source version of pywrap_tensorflow_macro under base name as well. pywrap_tensorflow_macro = pywrap_tensorflow_macro_opensource +# TODO(b/356020232): keep only the use_pywrap_rules part after migration is done +# also remove the comments below, as they will become false # This macro is for running python tests against system installed pip package # on Windows. # @@ -2547,23 +2614,45 @@ pywrap_tensorflow_macro = pywrap_tensorflow_macro_opensource # Note that this only works on Windows. See the definition of # //third_party/tensorflow/tools/pip_package:win_pip_package_marker for specific reasons. # 2. When --define=no_tensorflow_py_deps=false (by default), it's a normal py_test. -def py_test(deps = [], data = [], kernels = [], exec_properties = None, test_rule = _plain_py_test, **kwargs): +def py_test( + deps = [], + data = [], + kernels = [], + exec_properties = None, + test_rule = _plain_py_test, + env = {}, + **kwargs): if not exec_properties: exec_properties = tf_exec_properties(kwargs) - _make_tags_mutable(kwargs) - test_rule( - deps = select({ - "//conditions:default": deps, - clean_dep("//tensorflow:no_tensorflow_py_deps"): [], - }), - data = data + select({ - "//conditions:default": kernels, - clean_dep("//tensorflow:no_tensorflow_py_deps"): [], - }), - exec_properties = exec_properties, - **kwargs - ) + if use_pywrap_rules(): + test_env = { + "PYWRAP_TARGET": clean_dep(Label("//tensorflow/python:_pywrap_tensorflow")), + } + test_env.update(env) + actual_deps = deps.to_list() if hasattr(deps, "to_list") else deps + test_rule( + deps = actual_deps + [test_env["PYWRAP_TARGET"]], + exec_properties = exec_properties, + env = test_env, + data = data, + **kwargs + ) + else: + _make_tags_mutable(kwargs) + test_rule( + deps = select({ + "//conditions:default": deps, + clean_dep("//tensorflow:no_tensorflow_py_deps"): [], + }), + data = data + select({ + "//conditions:default": kernels, + clean_dep("//tensorflow:no_tensorflow_py_deps"): [], + }), + exec_properties = exec_properties, + env = env, + **kwargs + ) register_extension_info( extension = py_test, @@ -2596,11 +2685,15 @@ def pytype_library(name, pytype_deps = [], pytype_srcs = [], **kwargs): _make_tags_mutable(kwargs) _plain_py_library(name = name, **kwargs) +# TODO(b/356020232): remove completely after migration is done # Tensorflow uses rules_python 0.0.1, and in that version of rules_python, # the rules require the tags value to be a mutable list because they # modify it in-place. Later versions of rules_python don't have this # requirement. def _make_tags_mutable(kwargs): + if use_pywrap_rules(): + return + if "tags" in kwargs and kwargs["tags"] != None: # The value might be a frozen list, which looks just like # a regular list. So always make a copy. @@ -3021,6 +3114,7 @@ def pybind_library( **kwargs ) +# TODO(b/356020232): remove completely after migration is done # buildozer: disable=function-docstring-args def pybind_extension_opensource( name, @@ -3233,7 +3327,8 @@ def pybind_extension_opensource( ) # Export open source version of pybind_extension under base name as well. -pybind_extension = pybind_extension_opensource +pybind_extension = _pybind_extension if use_pywrap_rules() else pybind_extension_opensource +stripped_cc_info = _stripped_cc_info # Note: we cannot add //third_party/tf_runtime:__subpackages__ here, # because that builds all of tf_runtime's packages, and some of them @@ -3252,7 +3347,11 @@ def tsl_async_value_deps(): "@tf_runtime//third_party/llvm_derived:in_place", ] +# TODO(b/356020232): remove completely after migration is done def tf_python_pybind_static_deps(testonly = False): + if use_pywrap_rules(): + return [] + # TODO(b/146808376): Reduce the dependencies to those that are really needed. static_deps = [ "//:__subpackages__", @@ -3323,6 +3422,7 @@ def tf_python_pybind_static_deps(testonly = False): ] return if_oss(static_deps) +# TODO(b/356020232): remove completely after migration is done # buildozer: enable=function-docstring-args def tf_python_pybind_extension_opensource( name, @@ -3342,7 +3442,8 @@ def tf_python_pybind_extension_opensource( pytype_srcs = [], testonly = False, visibility = None, - win_def_file = None): + win_def_file = None, + additional_exported_symbols = None): """A wrapper macro for pybind_extension_opensource that is used in tensorflow/python/BUILD. Please do not use it anywhere else as it may behave unexpectedly. b/146445820 @@ -3374,7 +3475,7 @@ def tf_python_pybind_extension_opensource( ) # Export open source version of tf_python_pybind_extension under base name as well. -tf_python_pybind_extension = tf_python_pybind_extension_opensource +tf_python_pybind_extension = _pybind_extension if use_pywrap_rules() else tf_python_pybind_extension_opensource def tf_pybind_cc_library_wrapper_opensource(name, deps, visibility = None, **kwargs): """Wrapper for cc_library and proto dependencies used by tf_python_pybind_extension_opensource. @@ -3587,3 +3688,13 @@ def if_wheel_dependency(if_true, if_false = []): "@local_xla//xla/tsl:enable_wheel_dependency": if_true, "//conditions:default": if_false, }) + +# TODO(b/356020232): remove completely after migration is done +def pywrap_aware_tf_cc_shared_object(name, **kwargs): + if use_pywrap_rules(): + pass + else: + tf_cc_shared_object( + name = name, + **kwargs + ) diff --git a/tensorflow/tensorflow.default.bzl b/tensorflow/tensorflow.default.bzl index be61ba2e7b5..9f29115a44e 100644 --- a/tensorflow/tensorflow.default.bzl +++ b/tensorflow/tensorflow.default.bzl @@ -1,11 +1,21 @@ """Default (OSS) build versions of TensorFlow general-purpose build extensions.""" +load( + "@local_tsl//third_party/py/rules_pywrap:pywrap.bzl", + _pywrap_aware_cc_import = "pywrap_aware_cc_import", + _pywrap_aware_filegroup = "pywrap_aware_filegroup", + _pywrap_aware_genrule = "pywrap_aware_genrule", + _pywrap_common_library = "pywrap_common_library", + _pywrap_library = "pywrap_library", + _stripped_cc_info = "stripped_cc_info", +) load( "//tensorflow:tensorflow.bzl", _ADDITIONAL_API_INDEXABLE_SETTINGS = "ADDITIONAL_API_INDEXABLE_SETTINGS", _cc_header_only_library = "cc_header_only_library", _clean_dep = "clean_dep", _cuda_py_test = "cuda_py_test", + _custom_op_cc_header_only_library = "custom_op_cc_header_only_library", _filegroup = "filegroup", _genrule = "genrule", _get_compatible_with_portable = "get_compatible_with_portable", @@ -16,6 +26,7 @@ load( _pybind_extension = "pybind_extension", _pybind_library = "pybind_library", _pytype_library = "pytype_library", + _pywrap_aware_tf_cc_shared_object = "pywrap_aware_tf_cc_shared_object", _pywrap_tensorflow_macro = "pywrap_tensorflow_macro", _replace_with_portable_tf_lib_when_required = "replace_with_portable_tf_lib_when_required", _tensorflow_opensource_extra_deps = "tensorflow_opensource_extra_deps", @@ -80,6 +91,7 @@ tf_grpc_dependencies = _tf_grpc_dependencies tf_grpc_cc_dependencies = _tf_grpc_cc_dependencies get_compatible_with_portable = _get_compatible_with_portable cc_header_only_library = _cc_header_only_library +custom_op_cc_header_only_library = _custom_op_cc_header_only_library tf_gen_op_libs = _tf_gen_op_libs tf_gen_op_wrapper_cc = _tf_gen_op_wrapper_cc tf_gen_op_wrappers_cc = _tf_gen_op_wrappers_cc @@ -91,3 +103,10 @@ internal_tfrt_deps = _internal_tfrt_deps tf_disable_ptxas_warning_flags = _tf_disable_ptxas_warning_flags replace_with_portable_tf_lib_when_required = _replace_with_portable_tf_lib_when_required tf_python_framework_friends = _tf_python_framework_friends +pywrap_aware_tf_cc_shared_object = _pywrap_aware_tf_cc_shared_object +pywrap_aware_filegroup = _pywrap_aware_filegroup +pywrap_aware_genrule = _pywrap_aware_genrule +pywrap_aware_cc_import = _pywrap_aware_cc_import +pywrap_library = _pywrap_library +pywrap_common_library = _pywrap_common_library +stripped_cc_info = _stripped_cc_info diff --git a/tensorflow/tools/common/BUILD b/tensorflow/tools/common/BUILD index 6b994ebb7d9..75023a5d88e 100644 --- a/tensorflow/tools/common/BUILD +++ b/tensorflow/tools/common/BUILD @@ -3,6 +3,7 @@ load("//tensorflow:strict.default.bzl", "py_strict_library", "py_strict_test") load("//tensorflow:tensorflow.bzl", "VERSION_MAJOR") +load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap") package( # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], @@ -23,11 +24,13 @@ py_strict_library( py_strict_test( name = "public_api_test", srcs = ["public_api_test.py"], - data = select({ - "//tensorflow:macos": ["//tensorflow:libtensorflow_framework.dylib"], - "//tensorflow:windows": [], - "//conditions:default": ["//tensorflow:libtensorflow_framework.so.%s" % VERSION_MAJOR], - }), + data = if_pywrap( + if_false = select({ + "//tensorflow:macos": ["//tensorflow:libtensorflow_framework.dylib"], + "//tensorflow:windows": [], + "//conditions:default": ["//tensorflow:libtensorflow_framework.so.%s" % VERSION_MAJOR], + }), + ), python_version = "PY3", srcs_version = "PY3", deps = [ @@ -47,11 +50,13 @@ py_strict_library( py_strict_test( name = "traverse_test", srcs = ["traverse_test.py"], - data = select({ - "//tensorflow:macos": ["//tensorflow:libtensorflow_framework.dylib"], - "//tensorflow:windows": [], - "//conditions:default": ["//tensorflow:libtensorflow_framework.so.%s" % VERSION_MAJOR], - }), + data = if_pywrap( + if_false = select({ + "//tensorflow:macos": ["//tensorflow:libtensorflow_framework.dylib"], + "//tensorflow:windows": [], + "//conditions:default": ["//tensorflow:libtensorflow_framework.so.%s" % VERSION_MAJOR], + }), + ), python_version = "PY3", srcs_version = "PY3", deps = [ diff --git a/tensorflow/workspace2.bzl b/tensorflow/workspace2.bzl index be83c971749..779a1112ed5 100644 --- a/tensorflow/workspace2.bzl +++ b/tensorflow/workspace2.bzl @@ -797,7 +797,10 @@ def _tf_repositories(): urls = tf_mirror_urls("https://github.com/pybind/pybind11_protobuf/archive/80f3440cd8fee124e077e2e47a8a17b78b451363.zip"), sha256 = "c7ab64b1ccf9a678694a89035a8c865a693e4e872803778f91f0965c2f281d78", strip_prefix = "pybind11_protobuf-80f3440cd8fee124e077e2e47a8a17b78b451363", - patch_file = ["//third_party/pybind11_protobuf:remove_license.patch"], + patch_file = [ + "//third_party/pybind11_protobuf:protobuf.patch", + "//third_party/pybind11_protobuf:remove_license.patch", + ], ) tf_http_archive( diff --git a/third_party/protobuf/protobuf.patch b/third_party/protobuf/protobuf.patch index 9d928ba175f..ac33ccbf8c3 100644 --- a/third_party/protobuf/protobuf.patch +++ b/third_party/protobuf/protobuf.patch @@ -1,22 +1,46 @@ diff --git a/BUILD.bazel b/BUILD.bazel --- a/BUILD.bazel (revision 90b73ac3f0b10320315c2ca0d03a5a9b095d2f66) -+++ b/BUILD.bazel (date 1670471682469) -@@ -68,6 +68,7 @@ ++++ b/BUILD.bazel (date 1714620794503) +@@ -68,6 +68,8 @@ copts = COPTS, includes = ["src/"], linkopts = LINK_OPTS, ++ local_defines = ["PROTOBUF_USE_DLLS", "LIBPROTOBUF_EXPORTS"], + alwayslink = 1, visibility = ["//visibility:public"], ) -@@ -135,6 +136,7 @@ +@@ -135,6 +137,8 @@ copts = COPTS, includes = ["src/"], linkopts = LINK_OPTS, ++ local_defines = ["PROTOBUF_USE_DLLS", "LIBPROTOBUF_EXPORTS"], + alwayslink = 1, visibility = ["//visibility:public"], deps = [":protobuf_lite"] + select({ "//build_defs:config_msvc": [], +@@ -1074,7 +1078,8 @@ + "@com_google_protobuf//:type_proto", + "@com_google_protobuf//:wrappers_proto", + ], +- command_line = "--cpp_out=$(OUT)", ++ command_line = "--cpp_out=dllexport_decl=PROTOBUF_EXPORT:$(OUT)", ++# command_line = "--cpp_out=$(OUT)", + runtime = ":protobuf", + visibility = ["//visibility:public"], + ) +diff --git a/protobuf.bzl b/protobuf.bzl +--- a/protobuf.bzl (revision 90b73ac3f0b10320315c2ca0d03a5a9b095d2f66) ++++ b/protobuf.bzl (date 1714611573270) +@@ -127,7 +127,7 @@ + use_grpc_plugin = (ctx.attr.plugin_language == "grpc" and ctx.attr.plugin) + path_tpl = "$(realpath %s)" if in_gen_dir else "%s" + if ctx.attr.gen_cc: +- args += [("--cpp_out=" + path_tpl) % gen_dir] ++ args += [("--cpp_out=dllexport_decl=PROTOBUF_EXPORT:" + path_tpl) % gen_dir] + outs.extend(_CcOuts([src.basename], use_grpc_plugin = use_grpc_plugin)) + if ctx.attr.gen_py: + args += [("--python_out=" + path_tpl) % gen_dir] diff --git a/python/google/protobuf/pyext/descriptor.cc b/python/google/protobuf/pyext/descriptor.cc index 162531226..e93ec4809 100644 --- a/python/google/protobuf/pyext/descriptor.cc diff --git a/third_party/py/python_repo.bzl b/third_party/py/python_repo.bzl index 83778b744e0..ff1f4e8da21 100644 --- a/third_party/py/python_repo.bzl +++ b/third_party/py/python_repo.bzl @@ -53,6 +53,13 @@ Please check python_init_repositories() in your WORKSPACE file. merged_requirements_content, ) + use_pywrap_rules = bool( + ctx.os.environ.get("USE_PYWRAP_RULES", False), + ) + + if use_pywrap_rules: + print("!!!Using pywrap rules instead of directly creating .so objects!!!") + ctx.file( "py_version.bzl", """ @@ -62,12 +69,14 @@ WHEEL_NAME = "{wheel_name}" WHEEL_COLLAB = "{wheel_collab}" REQUIREMENTS = "{requirements}" REQUIREMENTS_WITH_LOCAL_WHEELS = "{requirements_with_local_wheels}" +USE_PYWRAP_RULES = {use_pywrap_rules} """.format( version = version, wheel_name = wheel_name, wheel_collab = wheel_collab, requirements = str(requirements), requirements_with_local_wheels = requirements_with_local_wheels, + use_pywrap_rules = use_pywrap_rules, ), ) @@ -185,6 +194,7 @@ python_repository = repository_rule( "HERMETIC_PYTHON_VERSION", "WHEEL_NAME", "WHEEL_COLLAB", + "USE_PYWRAP_RULES", ], local = True, ) diff --git a/third_party/pybind11_protobuf/protobuf.patch b/third_party/pybind11_protobuf/protobuf.patch new file mode 100644 index 00000000000..c568f5cd6f8 --- /dev/null +++ b/third_party/pybind11_protobuf/protobuf.patch @@ -0,0 +1,20 @@ +diff --git a/pybind11_protobuf/BUILD b/pybind11_protobuf/BUILD +--- a/pybind11_protobuf/BUILD (revision 80f3440cd8fee124e077e2e47a8a17b78b451363) ++++ b/pybind11_protobuf/BUILD (date 1714533560692) +@@ -53,8 +53,8 @@ + "proto_caster_impl.h", + ], + local_defines = select({ +- ":enable_pyproto_api_setting": ["PYBIND11_PROTOBUF_ENABLE_PYPROTO_API"], +- "//conditions:default": [], ++ ":enable_pyproto_api_setting": ["PROTOBUF_USE_DLLS", "PYBIND11_PROTOBUF_ENABLE_PYPROTO_API"], ++ "//conditions:default": ["PROTOBUF_USE_DLLS"], + }), + deps = [ + ":check_unknown_fields", +@@ -95,4 +95,5 @@ + "@com_google_absl//absl/synchronization", + "@com_google_protobuf//:protobuf", + ], ++ local_defines = ["PROTOBUF_USE_DLLS"], + ) \ No newline at end of file diff --git a/third_party/xla/third_party/py/non_hermetic/ml_dtypes/BUILD b/third_party/xla/third_party/py/non_hermetic/ml_dtypes/BUILD deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/third_party/xla/third_party/py/non_hermetic/ml_dtypes/LICENSE b/third_party/xla/third_party/py/non_hermetic/ml_dtypes/LICENSE deleted file mode 100644 index d6456956733..00000000000 --- a/third_party/xla/third_party/py/non_hermetic/ml_dtypes/LICENSE +++ /dev/null @@ -1,202 +0,0 @@ - - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/third_party/xla/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.BUILD b/third_party/xla/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.BUILD deleted file mode 100644 index f386124a36d..00000000000 --- a/third_party/xla/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.BUILD +++ /dev/null @@ -1,64 +0,0 @@ -load("@pybind11_bazel//:build_defs.bzl", "pybind_extension") - -package( - default_visibility = ["//visibility:public"], - licenses = ["notice"], -) - -exports_files(["LICENSE"]) - -cc_library( - name = "float8", - hdrs = ["include/float8.h"], - include_prefix = "ml_dtypes", - # Internal headers are all relative to . but other packages - # include these headers with the prefix. - includes = [ - ".", - "ml_dtypes", - ], - deps = ["@eigen_archive//:eigen3"], -) - -cc_library( - name = "intn", - hdrs = ["include/intn.h"], - include_prefix = "ml_dtypes", - # Internal headers are all relative to . but other packages - # include these headers with the prefix. - includes = [ - ".", - "ml_dtypes", - ], -) - -pybind_extension( - name = "_ml_dtypes_ext", - srcs = [ - "_src/common.h", - "_src/custom_float.h", - "_src/dtypes.cc", - "_src/int4_numpy.h", - "_src/numpy.cc", - "_src/numpy.h", - "_src/ufuncs.h", - ], - includes = ["ml_dtypes"], - visibility = [":__subpackages__"], - deps = [ - ":float8", - ":intn", - "@eigen_archive//:eigen3", - "@local_tsl//third_party/py/numpy:headers", - ], -) - -py_library( - name = "ml_dtypes", - srcs = [ - "__init__.py", - "_finfo.py", - "_iinfo.py", - ], - deps = [":_ml_dtypes_ext"], -) diff --git a/third_party/xla/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.tests.BUILD b/third_party/xla/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.tests.BUILD deleted file mode 100644 index c811379a19d..00000000000 --- a/third_party/xla/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.tests.BUILD +++ /dev/null @@ -1,71 +0,0 @@ -package( - default_visibility = ["//visibility:public"], -) - -py_library( - name = "testing_base", - deps = [ - "//:ml_dtypes", - "@absl_py//absl/testing:absltest", - "@absl_py//absl/testing:parameterized", - "@local_tsl//third_party/py/numpy", - ], -) - -py_test( - name = "custom_float_test", - srcs = ["custom_float_test.py"], - main = "custom_float_test.py", - deps = [":testing_base"], -) - -py_test( - name = "int4_test", - srcs = ["int4_test.py"], - main = "int4_test.py", - deps = [":testing_base"], -) - -py_test( - name = "iinfo_test", - srcs = ["iinfo_test.py"], - main = "iinfo_test.py", - deps = [":testing_base"], -) - -py_test( - name = "finfo_test", - srcs = ["finfo_test.py"], - main = "finfo_test.py", - deps = [":testing_base"], -) - -py_test( - name = "metadata_test", - srcs = ["metadata_test.py"], - main = "metadata_test.py", - deps = [":testing_base"], -) - -cc_test( - name = "float8_test", - srcs = ["float8_test.cc"], - linkstatic = 1, - deps = [ - "//:float8", - "@com_google_absl//absl/strings", - "@com_google_googletest//:gtest_main", - "@eigen_archive//:eigen3", - ], -) - -cc_test( - name = "intn_test_cc", - srcs = ["intn_test.cc"], - linkstatic = 1, - deps = [ - "//:intn", - "@com_google_googletest//:gtest_main", - "@eigen_archive//:eigen3", - ], -) diff --git a/third_party/xla/third_party/py/non_hermetic/ml_dtypes/workspace.bzl b/third_party/xla/third_party/py/non_hermetic/ml_dtypes/workspace.bzl deleted file mode 100644 index 51505bf3a14..00000000000 --- a/third_party/xla/third_party/py/non_hermetic/ml_dtypes/workspace.bzl +++ /dev/null @@ -1,22 +0,0 @@ -"""Provides the repo macro to import ml_dtypes. - -ml_dtypes provides machine-learning-specific data-types like bfloat16, -float8 varieties, and int4. -""" - -load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") - -def repo(): - ML_DTYPES_COMMIT = "24084d9ed2c3d45bf83b7a9bff833aa185bf9172" - ML_DTYPES_SHA256 = "c916a3e6b3d9bdcb476f506fdbbecb6d5e9f21f82f221dfcb42b320b4e85e55a" - tf_http_archive( - name = "ml_dtypes", - build_file = "//third_party/py/ml_dtypes:ml_dtypes.BUILD", - link_files = { - "//third_party/py/ml_dtypes:ml_dtypes.tests.BUILD": "tests/BUILD.bazel", - "//third_party/py/ml_dtypes:LICENSE": "LICENSE", - }, - sha256 = ML_DTYPES_SHA256, - strip_prefix = "ml_dtypes-{commit}/ml_dtypes".format(commit = ML_DTYPES_COMMIT), - urls = tf_mirror_urls("https://github.com/jax-ml/ml_dtypes/archive/{commit}/ml_dtypes-{commit}.tar.gz".format(commit = ML_DTYPES_COMMIT)), - ) diff --git a/third_party/xla/third_party/py/non_hermetic/numpy/BUILD b/third_party/xla/third_party/py/non_hermetic/numpy/BUILD deleted file mode 100644 index c80cc5287bc..00000000000 --- a/third_party/xla/third_party/py/non_hermetic/numpy/BUILD +++ /dev/null @@ -1,21 +0,0 @@ -licenses(["restricted"]) - -package(default_visibility = ["//visibility:public"]) - -py_library( - name = "numpy", - srcs = ["tf_numpy_dummy.py"], - srcs_version = "PY3", -) - -alias( - name = "headers", - actual = "@local_config_python//:numpy_headers", -) - -genrule( - name = "dummy", - outs = ["tf_numpy_dummy.py"], - cmd = "touch $@", - visibility = ["//visibility:private"], -) diff --git a/third_party/xla/third_party/py/non_hermetic/numpy/README.md b/third_party/xla/third_party/py/non_hermetic/numpy/README.md deleted file mode 100644 index 4e58b9df87b..00000000000 --- a/third_party/xla/third_party/py/non_hermetic/numpy/README.md +++ /dev/null @@ -1,4 +0,0 @@ -# numpy_ops - -The folder tf_numpy_api/ contains lists of NumPy API symbols that the -`numpy_ops` internal module in TensorFlow implements. diff --git a/third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/BUILD b/third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/BUILD deleted file mode 100644 index 070f8ab8a65..00000000000 --- a/third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/BUILD +++ /dev/null @@ -1,12 +0,0 @@ -# TensorFlow API backwards compatibility test goldens for tf.experimental.numpy. - -package( - # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], - default_visibility = ["//visibility:public"], - licenses = ["notice"], -) - -filegroup( - name = "api_golden", - srcs = glob(["*.pbtxt"]), -) diff --git a/third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.ndarray.pbtxt b/third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.ndarray.pbtxt deleted file mode 100644 index 9198264c029..00000000000 --- a/third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.ndarray.pbtxt +++ /dev/null @@ -1,51 +0,0 @@ -path: "tensorflow.experimental.numpy.ndarray" -tf_class { - is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" - member { - name: "OVERLOADABLE_OPERATORS" - mtype: "" - } - member { - name: "dtype" - mtype: "" - } - member { - name: "name" - mtype: "" - } - member { - name: "ndim" - mtype: "" - } - member { - name: "shape" - mtype: "" - } - member_method { - name: "__init__" - } - member_method { - name: "eval" - argspec: "args=[\'self\', \'feed_dict\', \'session\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "experimental_ref" - argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "get_shape" - argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "ref" - argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "set_shape" - argspec: "args=[\'self\', \'shape\'], varargs=None, keywords=None, defaults=None" - } -} diff --git a/third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.pbtxt b/third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.pbtxt deleted file mode 100644 index 2f5490ad0c9..00000000000 --- a/third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.pbtxt +++ /dev/null @@ -1,919 +0,0 @@ -path: "tensorflow.experimental.numpy" -tf_module { - member { - name: "bool_" - mtype: "" - } - member { - name: "complex128" - mtype: "" - } - member { - name: "complex64" - mtype: "" - } - member { - name: "complex_" - mtype: "" - } - member { - name: "e" - mtype: "" - } - member { - name: "float16" - mtype: "" - } - member { - name: "float32" - mtype: "" - } - member { - name: "float64" - mtype: "" - } - member { - name: "float_" - mtype: "" - } - member { - name: "iinfo" - mtype: "" - } - member { - name: "inexact" - mtype: "" - } - member { - name: "inf" - mtype: "" - } - member { - name: "int16" - mtype: "" - } - member { - name: "int32" - mtype: "" - } - member { - name: "int64" - mtype: "" - } - member { - name: "int8" - mtype: "" - } - member { - name: "int_" - mtype: "" - } - member { - name: "ndarray" - mtype: "" - } - member { - name: "newaxis" - mtype: "" - } - member { - name: "object_" - mtype: "" - } - member { - name: "pi" - mtype: "" - } - member { - name: "random" - mtype: "" - } - member { - name: "string_" - mtype: "" - } - member { - name: "uint16" - mtype: "" - } - member { - name: "uint32" - mtype: "" - } - member { - name: "uint64" - mtype: "" - } - member { - name: "uint8" - mtype: "" - } - member { - name: "unicode_" - mtype: "" - } - member_method { - name: "abs" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "absolute" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "add" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "all" - argspec: "args=[\'a\', \'axis\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "allclose" - argspec: "args=[\'a\', \'b\', \'rtol\', \'atol\', \'equal_nan\'], varargs=None, keywords=None, defaults=[\'1e-05\', \'1e-08\', \'False\'], " - } - member_method { - name: "amax" - argspec: "args=[\'a\', \'axis\', \'out\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " - } - member_method { - name: "amin" - argspec: "args=[\'a\', \'axis\', \'out\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " - } - member_method { - name: "angle" - argspec: "args=[\'z\', \'deg\'], varargs=None, keywords=None, defaults=[\'False\'], " - } - member_method { - name: "any" - argspec: "args=[\'a\', \'axis\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "append" - argspec: "args=[\'arr\', \'values\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "arange" - argspec: "args=[\'start\', \'stop\', \'step\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'None\'], " - } - member_method { - name: "arccos" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "arccosh" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "arcsin" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "arcsinh" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "arctan" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "arctan2" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "arctanh" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "argmax" - argspec: "args=[\'a\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "argmin" - argspec: "args=[\'a\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "argsort" - argspec: "args=[\'a\', \'axis\', \'kind\', \'order\'], varargs=None, keywords=None, defaults=[\'-1\', \'quicksort\', \'None\'], " - } - member_method { - name: "around" - argspec: "args=[\'a\', \'decimals\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "array" - argspec: "args=[\'val\', \'dtype\', \'copy\', \'ndmin\'], varargs=None, keywords=None, defaults=[\'None\', \'True\', \'0\'], " - } - member_method { - name: "array_equal" - argspec: "args=[\'a1\', \'a2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "asanyarray" - argspec: "args=[\'a\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "asarray" - argspec: "args=[\'a\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "ascontiguousarray" - argspec: "args=[\'a\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "atleast_1d" - argspec: "args=[], varargs=arys, keywords=None, defaults=None" - } - member_method { - name: "atleast_2d" - argspec: "args=[], varargs=arys, keywords=None, defaults=None" - } - member_method { - name: "atleast_3d" - argspec: "args=[], varargs=arys, keywords=None, defaults=None" - } - member_method { - name: "average" - argspec: "args=[\'a\', \'axis\', \'weights\', \'returned\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'False\'], " - } - member_method { - name: "bitwise_and" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "bitwise_not" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "bitwise_or" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "bitwise_xor" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "broadcast_arrays" - argspec: "args=[], varargs=args, keywords=kwargs, defaults=None" - } - member_method { - name: "broadcast_to" - argspec: "args=[\'array\', \'shape\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "cbrt" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "ceil" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "clip" - argspec: "args=[\'a\', \'a_min\', \'a_max\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "compress" - argspec: "args=[\'condition\', \'a\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "concatenate" - argspec: "args=[\'arys\', \'axis\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "conj" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "conjugate" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "copy" - argspec: "args=[\'a\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "cos" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "cosh" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "count_nonzero" - argspec: "args=[\'a\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "cross" - argspec: "args=[\'a\', \'b\', \'axisa\', \'axisb\', \'axisc\', \'axis\'], varargs=None, keywords=None, defaults=[\'-1\', \'-1\', \'-1\', \'None\'], " - } - member_method { - name: "cumprod" - argspec: "args=[\'a\', \'axis\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "cumsum" - argspec: "args=[\'a\', \'axis\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "deg2rad" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "diag" - argspec: "args=[\'v\', \'k\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "diag_indices" - argspec: "args=[\'n\', \'ndim\'], varargs=None, keywords=None, defaults=[\'2\'], " - } - member_method { - name: "diagflat" - argspec: "args=[\'v\', \'k\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "diagonal" - argspec: "args=[\'a\', \'offset\', \'axis1\', \'axis2\'], varargs=None, keywords=None, defaults=[\'0\', \'0\', \'1\'], " - } - member_method { - name: "diff" - argspec: "args=[\'a\', \'n\', \'axis\'], varargs=None, keywords=None, defaults=[\'1\', \'-1\'], " - } - member_method { - name: "divide" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "divmod" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "dot" - argspec: "args=[\'a\', \'b\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "dsplit" - argspec: "args=[\'ary\', \'indices_or_sections\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "dstack" - argspec: "args=[\'tup\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "einsum" - argspec: "args=[\'subscripts\'], varargs=operands, keywords=kwargs, defaults=None" - } - member_method { - name: "empty" - argspec: "args=[\'shape\', \'dtype\'], varargs=None, keywords=None, defaults=[\"\"], " - } - member_method { - name: "empty_like" - argspec: "args=[\'a\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "equal" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "exp" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "exp2" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "expand_dims" - argspec: "args=[\'a\', \'axis\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "experimental_enable_numpy_behavior" - argspec: "args=[\'prefer_float32\'], varargs=None, keywords=None, defaults=[\'False\'], " - } - member_method { - name: "expm1" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "eye" - argspec: "args=[\'N\', \'M\', \'k\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \"\"], " - } - member_method { - name: "fabs" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "finfo" - argspec: "args=[\'dtype\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "fix" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "flatten" - argspec: "args=[\'a\', \'order\'], varargs=None, keywords=None, defaults=[\'C\'], " - } - member_method { - name: "flip" - argspec: "args=[\'m\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "fliplr" - argspec: "args=[\'m\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "flipud" - argspec: "args=[\'m\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "float_power" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "floor" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "floor_divide" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "full" - argspec: "args=[\'shape\', \'fill_value\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "full_like" - argspec: "args=[\'a\', \'fill_value\', \'dtype\', \'order\', \'subok\', \'shape\'], varargs=None, keywords=None, defaults=[\'None\', \'K\', \'True\', \'None\'], " - } - member_method { - name: "gcd" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "geomspace" - argspec: "args=[\'start\', \'stop\', \'num\', \'endpoint\', \'dtype\', \'axis\'], varargs=None, keywords=None, defaults=[\'50\', \'True\', \'None\', \'0\'], " - } - member_method { - name: "greater" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "greater_equal" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "heaviside" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "hsplit" - argspec: "args=[\'ary\', \'indices_or_sections\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "hstack" - argspec: "args=[\'tup\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "hypot" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "identity" - argspec: "args=[\'n\', \'dtype\'], varargs=None, keywords=None, defaults=[\"\"], " - } - member_method { - name: "imag" - argspec: "args=[\'val\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "inner" - argspec: "args=[\'a\', \'b\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isclose" - argspec: "args=[\'a\', \'b\', \'rtol\', \'atol\', \'equal_nan\'], varargs=None, keywords=None, defaults=[\'1e-05\', \'1e-08\', \'False\'], " - } - member_method { - name: "iscomplex" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "iscomplexobj" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isfinite" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isinf" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isnan" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isneginf" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isposinf" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isreal" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isrealobj" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isscalar" - argspec: "args=[\'num\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "issubdtype" - argspec: "args=[\'arg1\', \'arg2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "ix_" - argspec: "args=[], varargs=args, keywords=None, defaults=None" - } - member_method { - name: "kron" - argspec: "args=[\'a\', \'b\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "lcm" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "less" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "less_equal" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "linspace" - argspec: "args=[\'start\', \'stop\', \'num\', \'endpoint\', \'retstep\', \'dtype\', \'axis\'], varargs=None, keywords=None, defaults=[\'50\', \'True\', \'False\', \"\", \'0\'], " - } - member_method { - name: "log" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "log10" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "log1p" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "log2" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "logaddexp" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "logaddexp2" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "logical_and" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "logical_not" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "logical_or" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "logical_xor" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "logspace" - argspec: "args=[\'start\', \'stop\', \'num\', \'endpoint\', \'base\', \'dtype\', \'axis\'], varargs=None, keywords=None, defaults=[\'50\', \'True\', \'10.0\', \'None\', \'0\'], " - } - member_method { - name: "matmul" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "max" - argspec: "args=[\'a\', \'axis\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "maximum" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "mean" - argspec: "args=[\'a\', \'axis\', \'dtype\', \'out\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], " - } - member_method { - name: "meshgrid" - argspec: "args=[], varargs=xi, keywords=kwargs, defaults=None" - } - member_method { - name: "min" - argspec: "args=[\'a\', \'axis\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "minimum" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "mod" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "moveaxis" - argspec: "args=[\'a\', \'source\', \'destination\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "multiply" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "nanmean" - argspec: "args=[\'a\', \'axis\', \'dtype\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " - } - member_method { - name: "nanprod" - argspec: "args=[\'a\', \'axis\', \'dtype\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'False\'], " - } - member_method { - name: "nansum" - argspec: "args=[\'a\', \'axis\', \'dtype\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'False\'], " - } - member_method { - name: "ndim" - argspec: "args=[\'a\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "negative" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "nextafter" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "nonzero" - argspec: "args=[\'a\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "not_equal" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "ones" - argspec: "args=[\'shape\', \'dtype\'], varargs=None, keywords=None, defaults=[\"\"], " - } - member_method { - name: "ones_like" - argspec: "args=[\'a\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "outer" - argspec: "args=[\'a\', \'b\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "pad" - argspec: "args=[\'array\', \'pad_width\', \'mode\'], varargs=None, keywords=kwargs, defaults=None" - } - member_method { - name: "polyval" - argspec: "args=[\'p\', \'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "positive" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "power" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "prod" - argspec: "args=[\'a\', \'axis\', \'dtype\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " - } - member_method { - name: "promote_types" - argspec: "args=[\'type1\', \'type2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "ptp" - argspec: "args=[\'a\', \'axis\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "rad2deg" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "ravel" - argspec: "args=[\'a\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "real" - argspec: "args=[\'val\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "reciprocal" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "remainder" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "repeat" - argspec: "args=[\'a\', \'repeats\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "reshape" - argspec: "args=[\'a\', \'newshape\', \'order\'], varargs=None, keywords=None, defaults=[\'C\'], " - } - member_method { - name: "result_type" - argspec: "args=[], varargs=arrays_and_dtypes, keywords=None, defaults=None" - } - member_method { - name: "roll" - argspec: "args=[\'a\', \'shift\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "rot90" - argspec: "args=[\'m\', \'k\', \'axes\'], varargs=None, keywords=None, defaults=[\'1\', \'(0, 1)\'], " - } - member_method { - name: "round" - argspec: "args=[\'a\', \'decimals\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "select" - argspec: "args=[\'condlist\', \'choicelist\', \'default\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "shape" - argspec: "args=[\'a\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "sign" - argspec: "args=[\'x\', \'out\', \'where\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\'], " - } - member_method { - name: "signbit" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "sin" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "sinc" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "sinh" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "size" - argspec: "args=[\'x\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "sort" - argspec: "args=[\'a\', \'axis\', \'kind\', \'order\'], varargs=None, keywords=None, defaults=[\'-1\', \'quicksort\', \'None\'], " - } - member_method { - name: "split" - argspec: "args=[\'ary\', \'indices_or_sections\', \'axis\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "sqrt" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "square" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "squeeze" - argspec: "args=[\'a\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "stack" - argspec: "args=[\'arrays\', \'axis\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "std" - argspec: "args=[\'a\', \'axis\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "subtract" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "sum" - argspec: "args=[\'a\', \'axis\', \'dtype\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " - } - member_method { - name: "swapaxes" - argspec: "args=[\'a\', \'axis1\', \'axis2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "take" - argspec: "args=[\'a\', \'indices\', \'axis\', \'out\', \'mode\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'clip\'], " - } - member_method { - name: "take_along_axis" - argspec: "args=[\'arr\', \'indices\', \'axis\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "tan" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "tanh" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "tensordot" - argspec: "args=[\'a\', \'b\', \'axes\'], varargs=None, keywords=None, defaults=[\'2\'], " - } - member_method { - name: "tile" - argspec: "args=[\'a\', \'reps\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "trace" - argspec: "args=[\'a\', \'offset\', \'axis1\', \'axis2\', \'dtype\'], varargs=None, keywords=None, defaults=[\'0\', \'0\', \'1\', \'None\'], " - } - member_method { - name: "transpose" - argspec: "args=[\'a\', \'axes\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "tri" - argspec: "args=[\'N\', \'M\', \'k\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\'], " - } - member_method { - name: "tril" - argspec: "args=[\'m\', \'k\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "triu" - argspec: "args=[\'m\', \'k\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "true_divide" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "vander" - argspec: "args=[\'x\', \'N\', \'increasing\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], " - } - member_method { - name: "var" - argspec: "args=[\'a\', \'axis\', \'dtype\', \'out\', \'ddof\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'0\', \'None\'], " - } - member_method { - name: "vdot" - argspec: "args=[\'a\', \'b\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "vsplit" - argspec: "args=[\'ary\', \'indices_or_sections\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "vstack" - argspec: "args=[\'tup\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "where" - argspec: "args=[\'condition\', \'x\', \'y\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "zeros" - argspec: "args=[\'shape\', \'dtype\'], varargs=None, keywords=None, defaults=[\"\"], " - } - member_method { - name: "zeros_like" - argspec: "args=[\'a\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\'], " - } -} diff --git a/third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.random.pbtxt b/third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.random.pbtxt deleted file mode 100644 index 61a4766f3f8..00000000000 --- a/third_party/xla/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.random.pbtxt +++ /dev/null @@ -1,35 +0,0 @@ -path: "tensorflow.experimental.numpy.random" -tf_module { - member_method { - name: "poisson" - argspec: "args=[\'lam\', \'size\'], varargs=None, keywords=None, defaults=[\'1.0\', \'None\'], " - } - member_method { - name: "rand" - argspec: "args=[], varargs=size, keywords=None, defaults=None" - } - member_method { - name: "randint" - argspec: "args=[\'low\', \'high\', \'size\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \"\"], " - } - member_method { - name: "randn" - argspec: "args=[], varargs=args, keywords=None, defaults=None" - } - member_method { - name: "random" - argspec: "args=[\'size\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "seed" - argspec: "args=[\'s\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "standard_normal" - argspec: "args=[\'size\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "uniform" - argspec: "args=[\'low\', \'high\', \'size\'], varargs=None, keywords=None, defaults=[\'0.0\', \'1.0\', \'None\'], " - } -} diff --git a/third_party/xla/third_party/py/python_repo.bzl b/third_party/xla/third_party/py/python_repo.bzl index 83778b744e0..ff1f4e8da21 100644 --- a/third_party/xla/third_party/py/python_repo.bzl +++ b/third_party/xla/third_party/py/python_repo.bzl @@ -53,6 +53,13 @@ Please check python_init_repositories() in your WORKSPACE file. merged_requirements_content, ) + use_pywrap_rules = bool( + ctx.os.environ.get("USE_PYWRAP_RULES", False), + ) + + if use_pywrap_rules: + print("!!!Using pywrap rules instead of directly creating .so objects!!!") + ctx.file( "py_version.bzl", """ @@ -62,12 +69,14 @@ WHEEL_NAME = "{wheel_name}" WHEEL_COLLAB = "{wheel_collab}" REQUIREMENTS = "{requirements}" REQUIREMENTS_WITH_LOCAL_WHEELS = "{requirements_with_local_wheels}" +USE_PYWRAP_RULES = {use_pywrap_rules} """.format( version = version, wheel_name = wheel_name, wheel_collab = wheel_collab, requirements = str(requirements), requirements_with_local_wheels = requirements_with_local_wheels, + use_pywrap_rules = use_pywrap_rules, ), ) @@ -185,6 +194,7 @@ python_repository = repository_rule( "HERMETIC_PYTHON_VERSION", "WHEEL_NAME", "WHEEL_COLLAB", + "USE_PYWRAP_RULES", ], local = True, ) diff --git a/third_party/xla/third_party/tsl/third_party/protobuf/protobuf.patch b/third_party/xla/third_party/tsl/third_party/protobuf/protobuf.patch index 9d928ba175f..ac33ccbf8c3 100644 --- a/third_party/xla/third_party/tsl/third_party/protobuf/protobuf.patch +++ b/third_party/xla/third_party/tsl/third_party/protobuf/protobuf.patch @@ -1,22 +1,46 @@ diff --git a/BUILD.bazel b/BUILD.bazel --- a/BUILD.bazel (revision 90b73ac3f0b10320315c2ca0d03a5a9b095d2f66) -+++ b/BUILD.bazel (date 1670471682469) -@@ -68,6 +68,7 @@ ++++ b/BUILD.bazel (date 1714620794503) +@@ -68,6 +68,8 @@ copts = COPTS, includes = ["src/"], linkopts = LINK_OPTS, ++ local_defines = ["PROTOBUF_USE_DLLS", "LIBPROTOBUF_EXPORTS"], + alwayslink = 1, visibility = ["//visibility:public"], ) -@@ -135,6 +136,7 @@ +@@ -135,6 +137,8 @@ copts = COPTS, includes = ["src/"], linkopts = LINK_OPTS, ++ local_defines = ["PROTOBUF_USE_DLLS", "LIBPROTOBUF_EXPORTS"], + alwayslink = 1, visibility = ["//visibility:public"], deps = [":protobuf_lite"] + select({ "//build_defs:config_msvc": [], +@@ -1074,7 +1078,8 @@ + "@com_google_protobuf//:type_proto", + "@com_google_protobuf//:wrappers_proto", + ], +- command_line = "--cpp_out=$(OUT)", ++ command_line = "--cpp_out=dllexport_decl=PROTOBUF_EXPORT:$(OUT)", ++# command_line = "--cpp_out=$(OUT)", + runtime = ":protobuf", + visibility = ["//visibility:public"], + ) +diff --git a/protobuf.bzl b/protobuf.bzl +--- a/protobuf.bzl (revision 90b73ac3f0b10320315c2ca0d03a5a9b095d2f66) ++++ b/protobuf.bzl (date 1714611573270) +@@ -127,7 +127,7 @@ + use_grpc_plugin = (ctx.attr.plugin_language == "grpc" and ctx.attr.plugin) + path_tpl = "$(realpath %s)" if in_gen_dir else "%s" + if ctx.attr.gen_cc: +- args += [("--cpp_out=" + path_tpl) % gen_dir] ++ args += [("--cpp_out=dllexport_decl=PROTOBUF_EXPORT:" + path_tpl) % gen_dir] + outs.extend(_CcOuts([src.basename], use_grpc_plugin = use_grpc_plugin)) + if ctx.attr.gen_py: + args += [("--python_out=" + path_tpl) % gen_dir] diff --git a/python/google/protobuf/pyext/descriptor.cc b/python/google/protobuf/pyext/descriptor.cc index 162531226..e93ec4809 100644 --- a/python/google/protobuf/pyext/descriptor.cc diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/BUILD b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/BUILD deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/LICENSE b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/LICENSE deleted file mode 100644 index d6456956733..00000000000 --- a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/LICENSE +++ /dev/null @@ -1,202 +0,0 @@ - - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.BUILD b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.BUILD deleted file mode 100644 index f386124a36d..00000000000 --- a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.BUILD +++ /dev/null @@ -1,64 +0,0 @@ -load("@pybind11_bazel//:build_defs.bzl", "pybind_extension") - -package( - default_visibility = ["//visibility:public"], - licenses = ["notice"], -) - -exports_files(["LICENSE"]) - -cc_library( - name = "float8", - hdrs = ["include/float8.h"], - include_prefix = "ml_dtypes", - # Internal headers are all relative to . but other packages - # include these headers with the prefix. - includes = [ - ".", - "ml_dtypes", - ], - deps = ["@eigen_archive//:eigen3"], -) - -cc_library( - name = "intn", - hdrs = ["include/intn.h"], - include_prefix = "ml_dtypes", - # Internal headers are all relative to . but other packages - # include these headers with the prefix. - includes = [ - ".", - "ml_dtypes", - ], -) - -pybind_extension( - name = "_ml_dtypes_ext", - srcs = [ - "_src/common.h", - "_src/custom_float.h", - "_src/dtypes.cc", - "_src/int4_numpy.h", - "_src/numpy.cc", - "_src/numpy.h", - "_src/ufuncs.h", - ], - includes = ["ml_dtypes"], - visibility = [":__subpackages__"], - deps = [ - ":float8", - ":intn", - "@eigen_archive//:eigen3", - "@local_tsl//third_party/py/numpy:headers", - ], -) - -py_library( - name = "ml_dtypes", - srcs = [ - "__init__.py", - "_finfo.py", - "_iinfo.py", - ], - deps = [":_ml_dtypes_ext"], -) diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.tests.BUILD b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.tests.BUILD deleted file mode 100644 index c811379a19d..00000000000 --- a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/ml_dtypes.tests.BUILD +++ /dev/null @@ -1,71 +0,0 @@ -package( - default_visibility = ["//visibility:public"], -) - -py_library( - name = "testing_base", - deps = [ - "//:ml_dtypes", - "@absl_py//absl/testing:absltest", - "@absl_py//absl/testing:parameterized", - "@local_tsl//third_party/py/numpy", - ], -) - -py_test( - name = "custom_float_test", - srcs = ["custom_float_test.py"], - main = "custom_float_test.py", - deps = [":testing_base"], -) - -py_test( - name = "int4_test", - srcs = ["int4_test.py"], - main = "int4_test.py", - deps = [":testing_base"], -) - -py_test( - name = "iinfo_test", - srcs = ["iinfo_test.py"], - main = "iinfo_test.py", - deps = [":testing_base"], -) - -py_test( - name = "finfo_test", - srcs = ["finfo_test.py"], - main = "finfo_test.py", - deps = [":testing_base"], -) - -py_test( - name = "metadata_test", - srcs = ["metadata_test.py"], - main = "metadata_test.py", - deps = [":testing_base"], -) - -cc_test( - name = "float8_test", - srcs = ["float8_test.cc"], - linkstatic = 1, - deps = [ - "//:float8", - "@com_google_absl//absl/strings", - "@com_google_googletest//:gtest_main", - "@eigen_archive//:eigen3", - ], -) - -cc_test( - name = "intn_test_cc", - srcs = ["intn_test.cc"], - linkstatic = 1, - deps = [ - "//:intn", - "@com_google_googletest//:gtest_main", - "@eigen_archive//:eigen3", - ], -) diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/workspace.bzl b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/workspace.bzl deleted file mode 100644 index 51505bf3a14..00000000000 --- a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/ml_dtypes/workspace.bzl +++ /dev/null @@ -1,22 +0,0 @@ -"""Provides the repo macro to import ml_dtypes. - -ml_dtypes provides machine-learning-specific data-types like bfloat16, -float8 varieties, and int4. -""" - -load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") - -def repo(): - ML_DTYPES_COMMIT = "24084d9ed2c3d45bf83b7a9bff833aa185bf9172" - ML_DTYPES_SHA256 = "c916a3e6b3d9bdcb476f506fdbbecb6d5e9f21f82f221dfcb42b320b4e85e55a" - tf_http_archive( - name = "ml_dtypes", - build_file = "//third_party/py/ml_dtypes:ml_dtypes.BUILD", - link_files = { - "//third_party/py/ml_dtypes:ml_dtypes.tests.BUILD": "tests/BUILD.bazel", - "//third_party/py/ml_dtypes:LICENSE": "LICENSE", - }, - sha256 = ML_DTYPES_SHA256, - strip_prefix = "ml_dtypes-{commit}/ml_dtypes".format(commit = ML_DTYPES_COMMIT), - urls = tf_mirror_urls("https://github.com/jax-ml/ml_dtypes/archive/{commit}/ml_dtypes-{commit}.tar.gz".format(commit = ML_DTYPES_COMMIT)), - ) diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/BUILD b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/BUILD deleted file mode 100644 index c80cc5287bc..00000000000 --- a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/BUILD +++ /dev/null @@ -1,21 +0,0 @@ -licenses(["restricted"]) - -package(default_visibility = ["//visibility:public"]) - -py_library( - name = "numpy", - srcs = ["tf_numpy_dummy.py"], - srcs_version = "PY3", -) - -alias( - name = "headers", - actual = "@local_config_python//:numpy_headers", -) - -genrule( - name = "dummy", - outs = ["tf_numpy_dummy.py"], - cmd = "touch $@", - visibility = ["//visibility:private"], -) diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/README.md b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/README.md deleted file mode 100644 index 4e58b9df87b..00000000000 --- a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/README.md +++ /dev/null @@ -1,4 +0,0 @@ -# numpy_ops - -The folder tf_numpy_api/ contains lists of NumPy API symbols that the -`numpy_ops` internal module in TensorFlow implements. diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/BUILD b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/BUILD deleted file mode 100644 index 070f8ab8a65..00000000000 --- a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/BUILD +++ /dev/null @@ -1,12 +0,0 @@ -# TensorFlow API backwards compatibility test goldens for tf.experimental.numpy. - -package( - # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], - default_visibility = ["//visibility:public"], - licenses = ["notice"], -) - -filegroup( - name = "api_golden", - srcs = glob(["*.pbtxt"]), -) diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.ndarray.pbtxt b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.ndarray.pbtxt deleted file mode 100644 index 9198264c029..00000000000 --- a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.ndarray.pbtxt +++ /dev/null @@ -1,51 +0,0 @@ -path: "tensorflow.experimental.numpy.ndarray" -tf_class { - is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" - is_instance: "" - member { - name: "OVERLOADABLE_OPERATORS" - mtype: "" - } - member { - name: "dtype" - mtype: "" - } - member { - name: "name" - mtype: "" - } - member { - name: "ndim" - mtype: "" - } - member { - name: "shape" - mtype: "" - } - member_method { - name: "__init__" - } - member_method { - name: "eval" - argspec: "args=[\'self\', \'feed_dict\', \'session\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "experimental_ref" - argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "get_shape" - argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "ref" - argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "set_shape" - argspec: "args=[\'self\', \'shape\'], varargs=None, keywords=None, defaults=None" - } -} diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.pbtxt b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.pbtxt deleted file mode 100644 index 2f5490ad0c9..00000000000 --- a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.pbtxt +++ /dev/null @@ -1,919 +0,0 @@ -path: "tensorflow.experimental.numpy" -tf_module { - member { - name: "bool_" - mtype: "" - } - member { - name: "complex128" - mtype: "" - } - member { - name: "complex64" - mtype: "" - } - member { - name: "complex_" - mtype: "" - } - member { - name: "e" - mtype: "" - } - member { - name: "float16" - mtype: "" - } - member { - name: "float32" - mtype: "" - } - member { - name: "float64" - mtype: "" - } - member { - name: "float_" - mtype: "" - } - member { - name: "iinfo" - mtype: "" - } - member { - name: "inexact" - mtype: "" - } - member { - name: "inf" - mtype: "" - } - member { - name: "int16" - mtype: "" - } - member { - name: "int32" - mtype: "" - } - member { - name: "int64" - mtype: "" - } - member { - name: "int8" - mtype: "" - } - member { - name: "int_" - mtype: "" - } - member { - name: "ndarray" - mtype: "" - } - member { - name: "newaxis" - mtype: "" - } - member { - name: "object_" - mtype: "" - } - member { - name: "pi" - mtype: "" - } - member { - name: "random" - mtype: "" - } - member { - name: "string_" - mtype: "" - } - member { - name: "uint16" - mtype: "" - } - member { - name: "uint32" - mtype: "" - } - member { - name: "uint64" - mtype: "" - } - member { - name: "uint8" - mtype: "" - } - member { - name: "unicode_" - mtype: "" - } - member_method { - name: "abs" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "absolute" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "add" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "all" - argspec: "args=[\'a\', \'axis\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "allclose" - argspec: "args=[\'a\', \'b\', \'rtol\', \'atol\', \'equal_nan\'], varargs=None, keywords=None, defaults=[\'1e-05\', \'1e-08\', \'False\'], " - } - member_method { - name: "amax" - argspec: "args=[\'a\', \'axis\', \'out\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " - } - member_method { - name: "amin" - argspec: "args=[\'a\', \'axis\', \'out\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " - } - member_method { - name: "angle" - argspec: "args=[\'z\', \'deg\'], varargs=None, keywords=None, defaults=[\'False\'], " - } - member_method { - name: "any" - argspec: "args=[\'a\', \'axis\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "append" - argspec: "args=[\'arr\', \'values\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "arange" - argspec: "args=[\'start\', \'stop\', \'step\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'1\', \'None\'], " - } - member_method { - name: "arccos" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "arccosh" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "arcsin" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "arcsinh" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "arctan" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "arctan2" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "arctanh" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "argmax" - argspec: "args=[\'a\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "argmin" - argspec: "args=[\'a\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "argsort" - argspec: "args=[\'a\', \'axis\', \'kind\', \'order\'], varargs=None, keywords=None, defaults=[\'-1\', \'quicksort\', \'None\'], " - } - member_method { - name: "around" - argspec: "args=[\'a\', \'decimals\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "array" - argspec: "args=[\'val\', \'dtype\', \'copy\', \'ndmin\'], varargs=None, keywords=None, defaults=[\'None\', \'True\', \'0\'], " - } - member_method { - name: "array_equal" - argspec: "args=[\'a1\', \'a2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "asanyarray" - argspec: "args=[\'a\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "asarray" - argspec: "args=[\'a\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "ascontiguousarray" - argspec: "args=[\'a\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "atleast_1d" - argspec: "args=[], varargs=arys, keywords=None, defaults=None" - } - member_method { - name: "atleast_2d" - argspec: "args=[], varargs=arys, keywords=None, defaults=None" - } - member_method { - name: "atleast_3d" - argspec: "args=[], varargs=arys, keywords=None, defaults=None" - } - member_method { - name: "average" - argspec: "args=[\'a\', \'axis\', \'weights\', \'returned\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'False\'], " - } - member_method { - name: "bitwise_and" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "bitwise_not" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "bitwise_or" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "bitwise_xor" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "broadcast_arrays" - argspec: "args=[], varargs=args, keywords=kwargs, defaults=None" - } - member_method { - name: "broadcast_to" - argspec: "args=[\'array\', \'shape\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "cbrt" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "ceil" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "clip" - argspec: "args=[\'a\', \'a_min\', \'a_max\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "compress" - argspec: "args=[\'condition\', \'a\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "concatenate" - argspec: "args=[\'arys\', \'axis\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "conj" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "conjugate" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "copy" - argspec: "args=[\'a\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "cos" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "cosh" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "count_nonzero" - argspec: "args=[\'a\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "cross" - argspec: "args=[\'a\', \'b\', \'axisa\', \'axisb\', \'axisc\', \'axis\'], varargs=None, keywords=None, defaults=[\'-1\', \'-1\', \'-1\', \'None\'], " - } - member_method { - name: "cumprod" - argspec: "args=[\'a\', \'axis\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "cumsum" - argspec: "args=[\'a\', \'axis\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "deg2rad" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "diag" - argspec: "args=[\'v\', \'k\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "diag_indices" - argspec: "args=[\'n\', \'ndim\'], varargs=None, keywords=None, defaults=[\'2\'], " - } - member_method { - name: "diagflat" - argspec: "args=[\'v\', \'k\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "diagonal" - argspec: "args=[\'a\', \'offset\', \'axis1\', \'axis2\'], varargs=None, keywords=None, defaults=[\'0\', \'0\', \'1\'], " - } - member_method { - name: "diff" - argspec: "args=[\'a\', \'n\', \'axis\'], varargs=None, keywords=None, defaults=[\'1\', \'-1\'], " - } - member_method { - name: "divide" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "divmod" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "dot" - argspec: "args=[\'a\', \'b\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "dsplit" - argspec: "args=[\'ary\', \'indices_or_sections\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "dstack" - argspec: "args=[\'tup\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "einsum" - argspec: "args=[\'subscripts\'], varargs=operands, keywords=kwargs, defaults=None" - } - member_method { - name: "empty" - argspec: "args=[\'shape\', \'dtype\'], varargs=None, keywords=None, defaults=[\"\"], " - } - member_method { - name: "empty_like" - argspec: "args=[\'a\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "equal" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "exp" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "exp2" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "expand_dims" - argspec: "args=[\'a\', \'axis\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "experimental_enable_numpy_behavior" - argspec: "args=[\'prefer_float32\'], varargs=None, keywords=None, defaults=[\'False\'], " - } - member_method { - name: "expm1" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "eye" - argspec: "args=[\'N\', \'M\', \'k\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \"\"], " - } - member_method { - name: "fabs" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "finfo" - argspec: "args=[\'dtype\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "fix" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "flatten" - argspec: "args=[\'a\', \'order\'], varargs=None, keywords=None, defaults=[\'C\'], " - } - member_method { - name: "flip" - argspec: "args=[\'m\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "fliplr" - argspec: "args=[\'m\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "flipud" - argspec: "args=[\'m\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "float_power" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "floor" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "floor_divide" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "full" - argspec: "args=[\'shape\', \'fill_value\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "full_like" - argspec: "args=[\'a\', \'fill_value\', \'dtype\', \'order\', \'subok\', \'shape\'], varargs=None, keywords=None, defaults=[\'None\', \'K\', \'True\', \'None\'], " - } - member_method { - name: "gcd" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "geomspace" - argspec: "args=[\'start\', \'stop\', \'num\', \'endpoint\', \'dtype\', \'axis\'], varargs=None, keywords=None, defaults=[\'50\', \'True\', \'None\', \'0\'], " - } - member_method { - name: "greater" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "greater_equal" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "heaviside" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "hsplit" - argspec: "args=[\'ary\', \'indices_or_sections\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "hstack" - argspec: "args=[\'tup\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "hypot" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "identity" - argspec: "args=[\'n\', \'dtype\'], varargs=None, keywords=None, defaults=[\"\"], " - } - member_method { - name: "imag" - argspec: "args=[\'val\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "inner" - argspec: "args=[\'a\', \'b\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isclose" - argspec: "args=[\'a\', \'b\', \'rtol\', \'atol\', \'equal_nan\'], varargs=None, keywords=None, defaults=[\'1e-05\', \'1e-08\', \'False\'], " - } - member_method { - name: "iscomplex" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "iscomplexobj" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isfinite" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isinf" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isnan" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isneginf" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isposinf" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isreal" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isrealobj" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "isscalar" - argspec: "args=[\'num\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "issubdtype" - argspec: "args=[\'arg1\', \'arg2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "ix_" - argspec: "args=[], varargs=args, keywords=None, defaults=None" - } - member_method { - name: "kron" - argspec: "args=[\'a\', \'b\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "lcm" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "less" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "less_equal" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "linspace" - argspec: "args=[\'start\', \'stop\', \'num\', \'endpoint\', \'retstep\', \'dtype\', \'axis\'], varargs=None, keywords=None, defaults=[\'50\', \'True\', \'False\', \"\", \'0\'], " - } - member_method { - name: "log" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "log10" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "log1p" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "log2" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "logaddexp" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "logaddexp2" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "logical_and" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "logical_not" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "logical_or" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "logical_xor" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "logspace" - argspec: "args=[\'start\', \'stop\', \'num\', \'endpoint\', \'base\', \'dtype\', \'axis\'], varargs=None, keywords=None, defaults=[\'50\', \'True\', \'10.0\', \'None\', \'0\'], " - } - member_method { - name: "matmul" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "max" - argspec: "args=[\'a\', \'axis\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "maximum" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "mean" - argspec: "args=[\'a\', \'axis\', \'dtype\', \'out\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], " - } - member_method { - name: "meshgrid" - argspec: "args=[], varargs=xi, keywords=kwargs, defaults=None" - } - member_method { - name: "min" - argspec: "args=[\'a\', \'axis\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "minimum" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "mod" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "moveaxis" - argspec: "args=[\'a\', \'source\', \'destination\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "multiply" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "nanmean" - argspec: "args=[\'a\', \'axis\', \'dtype\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " - } - member_method { - name: "nanprod" - argspec: "args=[\'a\', \'axis\', \'dtype\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'False\'], " - } - member_method { - name: "nansum" - argspec: "args=[\'a\', \'axis\', \'dtype\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'False\'], " - } - member_method { - name: "ndim" - argspec: "args=[\'a\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "negative" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "nextafter" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "nonzero" - argspec: "args=[\'a\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "not_equal" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "ones" - argspec: "args=[\'shape\', \'dtype\'], varargs=None, keywords=None, defaults=[\"\"], " - } - member_method { - name: "ones_like" - argspec: "args=[\'a\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "outer" - argspec: "args=[\'a\', \'b\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "pad" - argspec: "args=[\'array\', \'pad_width\', \'mode\'], varargs=None, keywords=kwargs, defaults=None" - } - member_method { - name: "polyval" - argspec: "args=[\'p\', \'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "positive" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "power" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "prod" - argspec: "args=[\'a\', \'axis\', \'dtype\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " - } - member_method { - name: "promote_types" - argspec: "args=[\'type1\', \'type2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "ptp" - argspec: "args=[\'a\', \'axis\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "rad2deg" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "ravel" - argspec: "args=[\'a\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "real" - argspec: "args=[\'val\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "reciprocal" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "remainder" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "repeat" - argspec: "args=[\'a\', \'repeats\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "reshape" - argspec: "args=[\'a\', \'newshape\', \'order\'], varargs=None, keywords=None, defaults=[\'C\'], " - } - member_method { - name: "result_type" - argspec: "args=[], varargs=arrays_and_dtypes, keywords=None, defaults=None" - } - member_method { - name: "roll" - argspec: "args=[\'a\', \'shift\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "rot90" - argspec: "args=[\'m\', \'k\', \'axes\'], varargs=None, keywords=None, defaults=[\'1\', \'(0, 1)\'], " - } - member_method { - name: "round" - argspec: "args=[\'a\', \'decimals\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "select" - argspec: "args=[\'condlist\', \'choicelist\', \'default\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "shape" - argspec: "args=[\'a\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "sign" - argspec: "args=[\'x\', \'out\', \'where\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\'], " - } - member_method { - name: "signbit" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "sin" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "sinc" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "sinh" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "size" - argspec: "args=[\'x\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "sort" - argspec: "args=[\'a\', \'axis\', \'kind\', \'order\'], varargs=None, keywords=None, defaults=[\'-1\', \'quicksort\', \'None\'], " - } - member_method { - name: "split" - argspec: "args=[\'ary\', \'indices_or_sections\', \'axis\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "sqrt" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "square" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "squeeze" - argspec: "args=[\'a\', \'axis\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "stack" - argspec: "args=[\'arrays\', \'axis\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "std" - argspec: "args=[\'a\', \'axis\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "subtract" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "sum" - argspec: "args=[\'a\', \'axis\', \'dtype\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], " - } - member_method { - name: "swapaxes" - argspec: "args=[\'a\', \'axis1\', \'axis2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "take" - argspec: "args=[\'a\', \'indices\', \'axis\', \'out\', \'mode\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'clip\'], " - } - member_method { - name: "take_along_axis" - argspec: "args=[\'arr\', \'indices\', \'axis\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "tan" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "tanh" - argspec: "args=[\'x\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "tensordot" - argspec: "args=[\'a\', \'b\', \'axes\'], varargs=None, keywords=None, defaults=[\'2\'], " - } - member_method { - name: "tile" - argspec: "args=[\'a\', \'reps\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "trace" - argspec: "args=[\'a\', \'offset\', \'axis1\', \'axis2\', \'dtype\'], varargs=None, keywords=None, defaults=[\'0\', \'0\', \'1\', \'None\'], " - } - member_method { - name: "transpose" - argspec: "args=[\'a\', \'axes\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "tri" - argspec: "args=[\'N\', \'M\', \'k\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'0\', \'None\'], " - } - member_method { - name: "tril" - argspec: "args=[\'m\', \'k\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "triu" - argspec: "args=[\'m\', \'k\'], varargs=None, keywords=None, defaults=[\'0\'], " - } - member_method { - name: "true_divide" - argspec: "args=[\'x1\', \'x2\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "vander" - argspec: "args=[\'x\', \'N\', \'increasing\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], " - } - member_method { - name: "var" - argspec: "args=[\'a\', \'axis\', \'dtype\', \'out\', \'ddof\', \'keepdims\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'0\', \'None\'], " - } - member_method { - name: "vdot" - argspec: "args=[\'a\', \'b\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "vsplit" - argspec: "args=[\'ary\', \'indices_or_sections\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "vstack" - argspec: "args=[\'tup\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "where" - argspec: "args=[\'condition\', \'x\', \'y\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " - } - member_method { - name: "zeros" - argspec: "args=[\'shape\', \'dtype\'], varargs=None, keywords=None, defaults=[\"\"], " - } - member_method { - name: "zeros_like" - argspec: "args=[\'a\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\'], " - } -} diff --git a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.random.pbtxt b/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.random.pbtxt deleted file mode 100644 index 61a4766f3f8..00000000000 --- a/third_party/xla/third_party/tsl/third_party/py/non_hermetic/numpy/tf_numpy_api/tensorflow.experimental.numpy.random.pbtxt +++ /dev/null @@ -1,35 +0,0 @@ -path: "tensorflow.experimental.numpy.random" -tf_module { - member_method { - name: "poisson" - argspec: "args=[\'lam\', \'size\'], varargs=None, keywords=None, defaults=[\'1.0\', \'None\'], " - } - member_method { - name: "rand" - argspec: "args=[], varargs=size, keywords=None, defaults=None" - } - member_method { - name: "randint" - argspec: "args=[\'low\', \'high\', \'size\', \'dtype\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \"\"], " - } - member_method { - name: "randn" - argspec: "args=[], varargs=args, keywords=None, defaults=None" - } - member_method { - name: "random" - argspec: "args=[\'size\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "seed" - argspec: "args=[\'s\'], varargs=None, keywords=None, defaults=None" - } - member_method { - name: "standard_normal" - argspec: "args=[\'size\'], varargs=None, keywords=None, defaults=[\'None\'], " - } - member_method { - name: "uniform" - argspec: "args=[\'low\', \'high\', \'size\'], varargs=None, keywords=None, defaults=[\'0.0\', \'1.0\', \'None\'], " - } -} diff --git a/third_party/xla/third_party/tsl/third_party/py/python_repo.bzl b/third_party/xla/third_party/tsl/third_party/py/python_repo.bzl index 83778b744e0..ff1f4e8da21 100644 --- a/third_party/xla/third_party/tsl/third_party/py/python_repo.bzl +++ b/third_party/xla/third_party/tsl/third_party/py/python_repo.bzl @@ -53,6 +53,13 @@ Please check python_init_repositories() in your WORKSPACE file. merged_requirements_content, ) + use_pywrap_rules = bool( + ctx.os.environ.get("USE_PYWRAP_RULES", False), + ) + + if use_pywrap_rules: + print("!!!Using pywrap rules instead of directly creating .so objects!!!") + ctx.file( "py_version.bzl", """ @@ -62,12 +69,14 @@ WHEEL_NAME = "{wheel_name}" WHEEL_COLLAB = "{wheel_collab}" REQUIREMENTS = "{requirements}" REQUIREMENTS_WITH_LOCAL_WHEELS = "{requirements_with_local_wheels}" +USE_PYWRAP_RULES = {use_pywrap_rules} """.format( version = version, wheel_name = wheel_name, wheel_collab = wheel_collab, requirements = str(requirements), requirements_with_local_wheels = requirements_with_local_wheels, + use_pywrap_rules = use_pywrap_rules, ), ) @@ -185,6 +194,7 @@ python_repository = repository_rule( "HERMETIC_PYTHON_VERSION", "WHEEL_NAME", "WHEEL_COLLAB", + "USE_PYWRAP_RULES", ], local = True, ) diff --git a/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/BUILD b/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/BUILD new file mode 100644 index 00000000000..595b43626f0 --- /dev/null +++ b/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/BUILD @@ -0,0 +1,16 @@ +load("@bazel_skylib//:bzl_library.bzl", "bzl_library") + +# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"]) + +exports_files(["pybind_extension.py.tpl"]) + +bzl_library( + name = "pywrap_bzl", + srcs = [ + "pywrap.bzl", + # copybara:uncomment "pywrap.google.bzl", + "pywrap.impl.bzl", + ], + # copybara:uncomment parse_tests = False, + visibility = ["//visibility:public"], +) diff --git a/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pybind_extension.py.tpl b/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pybind_extension.py.tpl new file mode 100644 index 00000000000..98428b51486 --- /dev/null +++ b/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pybind_extension.py.tpl @@ -0,0 +1,49 @@ +import os +import re + + +def __calc_import_path(): + module_name = os.path.basename(__file__)[:-3] + outer_module_name = "" # template_val + for var in ["PYWRAP_TARGET", "TEST_TARGET"]: + path = __find_pywrap_module_by_target_label(os.environ.get(var)) + if path: + return "%s.%s%s" % (path, outer_module_name, module_name) + + for var in ["RUNFILES_MANIFEST_FILE", "RUNFILES_DIR"]: + path = __find_pywrap_module_by_runfiles_env(os.environ.get(var)) + if path: + return "%s.%s%s" % (path, outer_module_name, module_name) + + raise RuntimeError("Could not detect original test/binary location") + + +def __find_pywrap_module_by_target_label(target_label): + if target_label: + return target_label.split("//", 1)[1].split(":")[0].replace("/", ".") + return None + + +def __find_pywrap_module_by_runfiles_env(runfiles_env_var): + pattern = re.compile( + r"bazel-out/.*/bin/(?P[\w/]*)/(?P\w+)(\.exe)?\.runfiles" + ) + if runfiles_env_var: + match = pattern.search(runfiles_env_var) + return match.group("pkg").replace("/", ".") + return None + + +def __update_globals(pywrap_m): + if hasattr(pywrap_m, '__all__'): + all_names = pywrap_m.__all__ + else: + all_names = [name for name in dir(pywrap_m) if not name.startswith('_')] + + extra_names = [] # template_val + all_names.extend(extra_names) + globals().update({name: getattr(pywrap_m, name) for name in all_names}) + + +__pywrap_m = __import__(__calc_import_path(), fromlist=["*"]) +__update_globals(__pywrap_m) diff --git a/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pywrap.bzl b/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pywrap.bzl new file mode 100644 index 00000000000..e7b038f571c --- /dev/null +++ b/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pywrap.bzl @@ -0,0 +1,20 @@ +load( + "//third_party/py/rules_pywrap:pywrap.default.bzl", + _pybind_extension = "pybind_extension", + _pywrap_aware_cc_import = "pywrap_aware_cc_import", + _pywrap_aware_filegroup = "pywrap_aware_filegroup", + _pywrap_aware_genrule = "pywrap_aware_genrule", + _pywrap_common_library = "pywrap_common_library", + _pywrap_library = "pywrap_library", + _stripped_cc_info = "stripped_cc_info", + _use_pywrap_rules = "use_pywrap_rules", +) + +pybind_extension = _pybind_extension +use_pywrap_rules = _use_pywrap_rules +pywrap_library = _pywrap_library +pywrap_common_library = _pywrap_common_library +stripped_cc_info = _stripped_cc_info +pywrap_aware_filegroup = _pywrap_aware_filegroup +pywrap_aware_genrule = _pywrap_aware_genrule +pywrap_aware_cc_import = _pywrap_aware_cc_import diff --git a/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pywrap.default.bzl b/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pywrap.default.bzl new file mode 100644 index 00000000000..b1514f100a4 --- /dev/null +++ b/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pywrap.default.bzl @@ -0,0 +1,146 @@ +# TODO(b/356020232): remove entire file and all usages after migration is done +load("@python_version_repo//:py_version.bzl", "USE_PYWRAP_RULES") +load( + "//third_party/py/rules_pywrap:pywrap.impl.bzl", + _pybind_extension = "pybind_extension", + _pywrap_common_library = "pywrap_common_library", + _pywrap_library = "pywrap_library", + _stripped_cc_info = "stripped_cc_info", +) + +def pybind_extension( + name, # original + deps, # original + srcs = [], # original + private_deps = [], # original + visibility = None, # original + win_def_file = None, # original + testonly = None, # original + compatible_with = None, # original + outer_module_name = "", # deprecate + additional_exported_symbols = [], + data = None, # original + # Garbage parameters, exist only to maingain backward compatibility for + # a while. Will be removed once migration is fully completed + + # To patch top-level deps lists in sophisticated cases + pywrap_ignored_deps_filter = ["@pybind11", "@pybind11//:pybind11"], + pywrap_private_deps_filter = [ + "@pybind11_abseil//pybind11_abseil:absl_casters", + "@pybind11_abseil//pybind11_abseil:import_status_module", + "@pybind11_abseil//pybind11_abseil:status_casters", + "@pybind11_protobuf//pybind11_protobuf:native_proto_caster", + ], + pytype_srcs = None, # alias for data + hdrs = [], # merge into sources + pytype_deps = None, # ignore? + ignore_link_in_framework = None, # ignore + dynamic_deps = [], # ignore + static_deps = [], # ignore + enable_stub_generation = None, # ignore + module_name = None, # ignore + link_in_framework = None, # ignore + additional_stubgen_deps = None, # ignore + **kwargs): + _ignore = [ + ignore_link_in_framework, + dynamic_deps, + static_deps, + enable_stub_generation, + module_name, + link_in_framework, + additional_stubgen_deps, + pytype_deps, + ] + + private_deps_filter_dict = {k: None for k in pywrap_private_deps_filter} + ignored_deps_filter_dict = {k: None for k in pywrap_ignored_deps_filter} + + actual_srcs = srcs + hdrs + + actual_data = data + if pytype_srcs: + data = pytype_srcs + + actual_deps = [] + actual_private_deps = [] + actual_default_deps = ["@pybind11//:pybind11"] + + if type(deps) == list: + for dep in deps: + if dep in ignored_deps_filter_dict: + continue + if dep in private_deps_filter_dict: + actual_private_deps.append(dep) + continue + actual_deps.append(dep) + else: + actual_deps = deps + actual_default_deps = [] + + _pybind_extension( + name = name, + deps = actual_deps, + srcs = actual_srcs, + private_deps = actual_private_deps, + visibility = visibility, + win_def_file = win_def_file, + testonly = testonly, + compatible_with = compatible_with, + outer_module_name = outer_module_name, + additional_exported_symbols = additional_exported_symbols, + data = actual_data, + default_deps = actual_default_deps, + **kwargs + ) + +def use_pywrap_rules(): + return USE_PYWRAP_RULES + +def pywrap_library(name, **kwargs): + if use_pywrap_rules(): + _pywrap_library( + name = name, + **kwargs + ) + +def pywrap_common_library(name, **kwargs): + if use_pywrap_rules(): + _pywrap_common_library( + name = name, + **kwargs + ) + +def stripped_cc_info(name, **kwargs): + if use_pywrap_rules(): + _stripped_cc_info( + name = name, + **kwargs + ) + +def pywrap_aware_filegroup(name, **kwargs): + if use_pywrap_rules(): + pass + else: + native.filegroup( + name = name, + **kwargs + ) + +def pywrap_aware_genrule(name, **kwargs): + if use_pywrap_rules(): + pass + else: + native.genrule( + name = name, + **kwargs + ) + +def pywrap_aware_cc_import(name, **kwargs): + if use_pywrap_rules(): + pass + else: + native.cc_import( + name = name, + **kwargs + ) diff --git a/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pywrap.impl.bzl b/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pywrap.impl.bzl new file mode 100644 index 00000000000..f33012c3876 --- /dev/null +++ b/third_party/xla/third_party/tsl/third_party/py/rules_pywrap/pywrap.impl.bzl @@ -0,0 +1,731 @@ +load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "find_cpp_toolchain", "use_cpp_toolchain") + +PywrapInfo = provider( + fields = { + "cc_info": "Wrapped CcInfo", + "private_deps": "Libraries to link only to individual pywrap libraries, but not in commmon library", + "owner": "Owner's label", + "py_stub": "Pybind Python stub used to resolve cross-package references", + "outer_module_name": "Outer module name for deduping libraries with the same name", + "cc_only": "True if this PywrapInfo represents cc-only library (no PyIni_)", + }, +) + +CollectedPywrapInfo = provider( + fields = { + "pywrap_infos": "depset of PywrapInfo providers", + }, +) + +PywrapFilters = provider( + fields = { + "py_cc_linker_inputs": "", + "cc_linker_inputs": "", + "pywrap_private_linker_inputs": "", + }, +) + +def pywrap_library( + name, + deps, + py_cc_deps_filter = [], + cc_deps_filter = [], + linkopts = [], + py_cc_linkopts = [], + win_def_file = None, + py_cc_win_def_file = None, + pywrap_count = None, + extra_deps = ["@pybind11//:pybind11"], + visibility = None, + testonly = None, + compatible_with = None): + # 0) If pywrap_count is not specified, assume we pass pybind_extension, + # targets directly, so actual pywrap_count should just be equal to number + # of deps. + actual_pywrap_count = len(deps) if pywrap_count == None else pywrap_count + + # 1) Create common libraries cc-only (C API) and py-specific (parts reused + # by different pywrap libraries but dependin on Python symbols). + # The common library should link in everything except the object file with + # Python Extension's init function PyInit_. + info_collector_name = "_%s_info_collector" % name + collected_pywrap_infos( + name = info_collector_name, + deps = deps, + pywrap_count = actual_pywrap_count, + ) + + linker_input_filters_name = "_%s_linker_input_filters" % name + _linker_input_filters( + name = linker_input_filters_name, + dep = ":%s" % info_collector_name, + py_cc_deps_filter = py_cc_deps_filter, + cc_deps_filter = cc_deps_filter, + ) + + # _internal binary + common_split_name = "_%s_split" % name + _pywrap_split_library( + name = common_split_name, + mode = "cc_common", + dep = ":%s" % info_collector_name, + linker_input_filters = "%s" % linker_input_filters_name, + testonly = testonly, + compatible_with = compatible_with, + ) + + common_cc_binary_name = "%s_internal" % name + common_import_name = _construct_common_binary( + common_cc_binary_name, + [":%s" % common_split_name], + linkopts, + testonly, + compatible_with, + win_def_file, + None, + ) + + # _py_internal binary + py_common_split_name = "_%s_py_split" % name + _pywrap_split_library( + name = py_common_split_name, + mode = "py_common", + dep = ":%s" % info_collector_name, + linker_input_filters = "%s" % linker_input_filters_name, + testonly = testonly, + compatible_with = compatible_with, + ) + + common_py_cc_binary_name = "%s_py_internal" % name + common_py_import_name = _construct_common_binary( + common_py_cc_binary_name, + [ + ":%s" % py_common_split_name, + ":%s" % common_import_name, + "@pybind11//:pybind11", + ], + py_cc_linkopts, + testonly, + compatible_with, + py_cc_win_def_file, + ["PROTOBUF_USE_DLLS"], + ) + + common_deps = extra_deps + [ + ":%s" % common_import_name, + ":%s" % common_py_import_name, + ] + binaries_data = [ + ":%s" % common_cc_binary_name, + ":%s" % common_py_cc_binary_name, + ] + + # 2) Create individual super-thin pywrap libraries, which depend on the + # common one. The individual libraries must link in statically only the + # object file with Python Extension's init function PyInit_ + # + shared_objects = [] + for pywrap_index in range(0, actual_pywrap_count): + dep_name = "_%s_%s" % (name, pywrap_index) + shared_object_name = "%s_shared_object" % dep_name + win_def_name = "%s_win_def" % dep_name + pywrap_name = "%s_pywrap" % dep_name + + _pywrap_split_library( + name = pywrap_name, + mode = "pywrap", + dep = ":%s" % info_collector_name, + linker_input_filters = "%s" % linker_input_filters_name, + pywrap_index = pywrap_index, + testonly = testonly, + compatible_with = compatible_with, + ) + + _generated_win_def_file( + name = win_def_name, + dep = ":%s" % info_collector_name, + pywrap_index = pywrap_index, + testonly = testonly, + compatible_with = compatible_with, + ) + + native.cc_binary( + name = shared_object_name, + srcs = [], + deps = [":%s" % pywrap_name] + common_deps, + linkshared = True, + linkstatic = True, + win_def_file = ":%s" % win_def_name, + testonly = testonly, + compatible_with = compatible_with, + local_defines = ["PROTOBUF_USE_DLLS"], + ) + shared_objects.append(":%s" % shared_object_name) + + # 3) Construct final binaries with proper names and put them as data + # attribute in a py_library, which is the final and only public artifact of + # this macro + # + pywrap_binaries_name = "_%s_binaries" % name + _pywrap_binaries( + name = pywrap_binaries_name, + collected_pywraps = ":%s" % info_collector_name, + deps = shared_objects, + extension = select({ + "@bazel_tools//src/conditions:windows": ".pyd", + "//conditions:default": ".so", + }), + testonly = testonly, + compatible_with = compatible_with, + ) + + binaries_data.append("%s" % pywrap_binaries_name) + binaries_data.extend([shared_objects[0]]) + + native.py_library( + name = name, + srcs = [":%s" % info_collector_name], + data = binaries_data, + testonly = testonly, + compatible_with = compatible_with, + visibility = visibility, + ) + + # For debugging purposes only + native.filegroup( + name = "_%s_all_binaries" % name, + srcs = binaries_data, + testonly = testonly, + compatible_with = compatible_with, + ) + +def _construct_common_binary( + name, + deps, + linkopts, + testonly, + compatible_with, + win_def_file, + local_defines): + native.cc_binary( + name = name, + deps = deps, + linkstatic = True, + linkshared = True, + linkopts = linkopts, + testonly = testonly, + compatible_with = compatible_with, + win_def_file = win_def_file, + local_defines = local_defines, + ) + + if_lib_name = "%s_if_lib" % name + native.filegroup( + name = if_lib_name, + srcs = [":%s" % name], + output_group = "interface_library", + testonly = testonly, + compatible_with = compatible_with, + ) + + import_name = "%s_import" % name + native.cc_import( + name = import_name, + shared_library = ":%s" % name, + interface_library = ":%s" % if_lib_name, + testonly = testonly, + compatible_with = compatible_with, + ) + + return import_name + +def _pywrap_split_library_impl(ctx): + pywrap_index = ctx.attr.pywrap_index + pywrap_infos = ctx.attr.dep[CollectedPywrapInfo].pywrap_infos.to_list() + split_linker_inputs = [] + private_linker_inputs = [] + + mode = ctx.attr.mode + filters = ctx.attr.linker_input_filters[PywrapFilters] + py_cc_linker_inputs = filters.py_cc_linker_inputs + + if mode == "pywrap": + pw = pywrap_infos[pywrap_index] + + # print("%s matches %s" % (str(pw.owner), ctx.label)) + if not pw.cc_only: + li = pw.cc_info.linking_context.linker_inputs.to_list()[0] + split_linker_inputs.append(li) + private_linker_inputs = [ + depset(direct = filters.pywrap_private_linker_inputs[pywrap_index].keys()), + ] + else: + for i in range(0, len(pywrap_infos)): + pw = pywrap_infos[i] + pw_private_linker_inputs = filters.pywrap_private_linker_inputs[i] + pw_lis = pw.cc_info.linking_context.linker_inputs.to_list()[1:] + for li in pw_lis: + if li in pw_private_linker_inputs: + continue + if li in filters.py_cc_linker_inputs: + if mode == "py_common": + split_linker_inputs.append(li) + elif mode == "cc_common": + split_linker_inputs.append(li) + + dependency_libraries = _construct_dependency_libraries( + ctx, + split_linker_inputs, + ) + + linker_input = cc_common.create_linker_input( + owner = ctx.label, + libraries = depset(direct = dependency_libraries), + ) + + linking_context = cc_common.create_linking_context( + linker_inputs = depset( + direct = [linker_input], + transitive = private_linker_inputs, + ), + ) + + return [CcInfo(linking_context = linking_context)] + +_pywrap_split_library = rule( + attrs = { + "dep": attr.label( + allow_files = False, + providers = [CollectedPywrapInfo], + ), + # py_deps, meaning C++ deps which depend on Python symbols + "linker_input_filters": attr.label( + allow_files = False, + providers = [PywrapFilters], + mandatory = True, + ), + "pywrap_index": attr.int(mandatory = False, default = -1), + "mode": attr.string( + mandatory = True, + values = ["pywrap", "cc_common", "py_common"], + ), + "_cc_toolchain": attr.label( + default = "@bazel_tools//tools/cpp:current_cc_toolchain", + ), + }, + fragments = ["cpp"], + toolchains = use_cpp_toolchain(), + implementation = _pywrap_split_library_impl, +) + +def _construct_dependency_libraries(ctx, split_linker_inputs): + cc_toolchain = find_cpp_toolchain(ctx) + feature_configuration = cc_common.configure_features( + ctx = ctx, + cc_toolchain = cc_toolchain, + requested_features = ctx.features, + unsupported_features = ctx.disabled_features, + ) + dependency_libraries = [] + for split_linker_input in split_linker_inputs: + for lib in split_linker_input.libraries: + lib_copy = lib + if not lib.alwayslink: + lib_copy = cc_common.create_library_to_link( + actions = ctx.actions, + cc_toolchain = cc_toolchain, + feature_configuration = feature_configuration, + static_library = lib.static_library, + pic_static_library = lib.pic_static_library, + interface_library = lib.interface_library, + alwayslink = True, + ) + dependency_libraries.append(lib_copy) + + return dependency_libraries + +def _linker_input_filters_impl(ctx): + py_cc_linker_inputs = {} + for py_cc_dep in ctx.attr.py_cc_deps_filter: + for li in py_cc_dep[CcInfo].linking_context.linker_inputs.to_list()[:1]: + py_cc_linker_inputs[li] = li.owner + + cc_linker_inputs = {} + for cc_dep in ctx.attr.cc_deps_filter: + for li in cc_dep[CcInfo].linking_context.linker_inputs.to_list()[:1]: + cc_linker_inputs[li] = li.owner + + pywrap_infos = ctx.attr.dep[CollectedPywrapInfo].pywrap_infos.to_list() + pywrap_private_linker_inputs = [] + + for pw in pywrap_infos: + private_linker_inputs = {} + + for private_dep in pw.private_deps: + for priv_li in private_dep[CcInfo].linking_context.linker_inputs.to_list(): + if (priv_li not in py_cc_linker_inputs) and (priv_li not in cc_linker_inputs): + private_linker_inputs[priv_li] = priv_li.owner + pywrap_private_linker_inputs.append(private_linker_inputs) + + return [ + PywrapFilters( + py_cc_linker_inputs = py_cc_linker_inputs, + pywrap_private_linker_inputs = pywrap_private_linker_inputs, + ), + ] + +_linker_input_filters = rule( + attrs = { + "dep": attr.label( + allow_files = False, + providers = [CollectedPywrapInfo], + ), + "py_cc_deps_filter": attr.label_list( + allow_files = False, + providers = [CcInfo], + mandatory = False, + default = [], + ), + "cc_deps_filter": attr.label_list( + allow_files = False, + providers = [CcInfo], + mandatory = False, + default = [], + ), + }, + implementation = _linker_input_filters_impl, +) + +def pywrap_common_library(name, dep): + native.alias( + name = name, + actual = "%s_internal_import" % dep, + ) + +def pywrap_py_common_library(name, dep): + native.alias( + name = name, + actual = "%s_py_internal_import" % dep, + ) + +def _generated_win_def_file_impl(ctx): + pywrap_infos = ctx.attr.dep[CollectedPywrapInfo].pywrap_infos.to_list() + pywrap_info = pywrap_infos[ctx.attr.pywrap_index] + win_def_file_name = "%s.def" % pywrap_info.owner.name + win_def_file = ctx.actions.declare_file(win_def_file_name) + + if pywrap_info.cc_only: + command = "echo \"EXPORTS\r\n\">> {win_def_file}" + else: + command = "echo \"EXPORTS\r\n PyInit_{owner}\">> {win_def_file}" + + ctx.actions.run_shell( + inputs = [], + command = command.format( + owner = pywrap_info.owner.name, + win_def_file = win_def_file.path, + ), + outputs = [win_def_file], + ) + + return [DefaultInfo(files = depset(direct = [win_def_file]))] + +_generated_win_def_file = rule( + attrs = { + "dep": attr.label( + allow_files = False, + providers = [CollectedPywrapInfo], + ), + "pywrap_index": attr.int(mandatory = True), + }, + implementation = _generated_win_def_file_impl, +) + +def pybind_extension( + name, + deps, + srcs = [], + private_deps = [], + visibility = None, + win_def_file = None, + testonly = None, + compatible_with = None, + outer_module_name = "", + additional_exported_symbols = [], + default_deps = ["@pybind11//:pybind11"], + **kwargs): + cc_library_name = "_%s_cc_library" % name + + native.cc_library( + name = cc_library_name, + deps = deps + private_deps + default_deps, + srcs = srcs, + linkstatic = True, + alwayslink = True, + visibility = visibility, + testonly = testonly, + compatible_with = compatible_with, + local_defines = ["PROTOBUF_USE_DLLS"], + **kwargs + ) + + if not srcs: + _cc_only_pywrap_info_wrapper( + name = name, + deps = ["%s" % cc_library_name], + testonly = testonly, + compatible_with = compatible_with, + visibility = visibility, + ) + else: + _pywrap_info_wrapper( + name = name, + deps = ["%s" % cc_library_name], + private_deps = private_deps, + outer_module_name = outer_module_name, + additional_exported_symbols = additional_exported_symbols, + testonly = testonly, + compatible_with = compatible_with, + visibility = visibility, + ) + +def _pywrap_info_wrapper_impl(ctx): + #the attribute is called deps not dep to match aspect's attr_aspects + + if len(ctx.attr.deps) != 1: + fail("deps attribute must contain exactly one dependency") + + py_stub = ctx.actions.declare_file("%s.py" % ctx.attr.name) + substitutions = {} + outer_module_name = ctx.attr.outer_module_name + if outer_module_name: + val = 'outer_module_name = "%s."' % outer_module_name + substitutions['outer_module_name = "" # template_val'] = val + + additional_exported_symbols = ctx.attr.additional_exported_symbols + if additional_exported_symbols: + val = "extra_names = %s # template_val" % additional_exported_symbols + substitutions["extra_names = [] # template_val"] = val + + ctx.actions.expand_template( + template = ctx.file.py_stub_src, + output = py_stub, + substitutions = substitutions, + ) + + return [ + PyInfo(transitive_sources = depset()), + PywrapInfo( + cc_info = ctx.attr.deps[0][CcInfo], + private_deps = ctx.attr.private_deps, + owner = ctx.label, + py_stub = py_stub, + outer_module_name = outer_module_name, + cc_only = False, + ), + ] + +_pywrap_info_wrapper = rule( + attrs = { + "deps": attr.label_list(providers = [CcInfo]), + "private_deps": attr.label_list(providers = [CcInfo]), + "outer_module_name": attr.string(mandatory = False, default = ""), + "py_stub_src": attr.label( + allow_single_file = True, + default = Label("//third_party/py/rules_pywrap:pybind_extension.py.tpl"), + ), + "additional_exported_symbols": attr.string_list( + mandatory = False, + default = [], + ), + }, + implementation = _pywrap_info_wrapper_impl, +) + +def _cc_only_pywrap_info_wrapper_impl(ctx): + wrapped_dep = ctx.attr.deps[0] + return [ + PyInfo(transitive_sources = depset()), + PywrapInfo( + cc_info = wrapped_dep[CcInfo], + private_deps = [], + owner = ctx.label, + py_stub = None, + outer_module_name = None, + cc_only = True, + ), + ] + +_cc_only_pywrap_info_wrapper = rule( + attrs = { + "deps": attr.label_list(providers = [CcInfo]), + }, + implementation = _cc_only_pywrap_info_wrapper_impl, +) + +def _pywrap_info_collector_aspect_impl(target, ctx): + pywrap_infos = [] + transitive_pywrap_infos = [] + if PywrapInfo in target: + pywrap_infos.append(target[PywrapInfo]) + + if hasattr(ctx.rule.attr, "deps"): + for dep in ctx.rule.attr.deps: + if CollectedPywrapInfo in dep: + collected_pywrap_info = dep[CollectedPywrapInfo] + transitive_pywrap_infos.append(collected_pywrap_info.pywrap_infos) + + return [ + CollectedPywrapInfo( + pywrap_infos = depset( + direct = pywrap_infos, + transitive = transitive_pywrap_infos, + order = "topological", + ), + ), + ] + +_pywrap_info_collector_aspect = aspect( + attr_aspects = ["deps"], + implementation = _pywrap_info_collector_aspect_impl, +) + +def _collected_pywrap_infos_impl(ctx): + pywrap_infos = [] + for dep in ctx.attr.deps: + if CollectedPywrapInfo in dep: + pywrap_infos.append(dep[CollectedPywrapInfo].pywrap_infos) + + rv = CollectedPywrapInfo( + pywrap_infos = depset( + transitive = pywrap_infos, + order = "topological", + ), + ) + pywraps = rv.pywrap_infos.to_list() + + if ctx.attr.pywrap_count != len(pywraps): + found_pywraps = "\n ".join([str(pw.owner) for pw in pywraps]) + fail(""" + Number of actual pywrap libraries does not match expected pywrap_count. + Expected pywrap_count: {expected_pywrap_count} + Actual pywrap_count: {actual_pywra_count} + Actual pywrap libraries in the transitive closure of {label}: + {found_pywraps} + """.format( + expected_pywrap_count = ctx.attr.pywrap_count, + actual_pywra_count = len(pywraps), + label = ctx.label, + found_pywraps = found_pywraps, + )) + + py_stubs = [] + for pw in pywraps: + if pw.py_stub: + py_stubs.append(pw.py_stub) + + return [ + DefaultInfo(files = depset(direct = py_stubs)), + rv, + ] + +collected_pywrap_infos = rule( + attrs = { + "deps": attr.label_list( + aspects = [_pywrap_info_collector_aspect], + providers = [PyInfo], + ), + "pywrap_count": attr.int(mandatory = True, default = 1), + }, + implementation = _collected_pywrap_infos_impl, +) + +def _pywrap_binaries_impl(ctx): + deps = ctx.attr.deps + dep = ctx.attr.collected_pywraps + extension = ctx.attr.extension + + pywrap_infos = dep[CollectedPywrapInfo].pywrap_infos.to_list() + original_binaries = deps + + if len(pywrap_infos) != len(original_binaries): + fail() + + final_binaries = [] + original_to_final_binaries = [ + "\n\nvvv Shared objects corresondence map, target = {} vvv".format(ctx.label), + ] + for i in range(0, len(pywrap_infos)): + pywrap_info = pywrap_infos[i] + original_binary = original_binaries[i] + subfolder = "" + if pywrap_info.outer_module_name: + subfolder = pywrap_info.outer_module_name + "/" + final_binary_name = "%s%s%s" % (subfolder, pywrap_info.owner.name, extension) + final_binary = ctx.actions.declare_file(final_binary_name) + original_binary_file = original_binary.files.to_list()[0] + ctx.actions.run_shell( + inputs = [original_binary_file], + command = "cp {original} {final}".format( + original = original_binary_file.path, + final = final_binary.path, + ), + outputs = [final_binary], + ) + + original_to_final_binaries.append( + " '{original}' => '{final}'".format( + original = original_binary_file.path, + final = final_binary.path, + ), + ) + + final_binaries.append(final_binary) + + original_to_final_binaries.append( + "^^^ Shared objects corresondence map^^^\n\n", + ) + print("\n".join(original_to_final_binaries)) + + return [DefaultInfo(files = depset(direct = final_binaries))] + +_pywrap_binaries = rule( + attrs = { + "deps": attr.label_list(mandatory = True, allow_files = False), + "collected_pywraps": attr.label(mandatory = True, allow_files = False), + "extension": attr.string(default = ".so"), + }, + implementation = _pywrap_binaries_impl, +) + +def _stripped_cc_info_impl(ctx): + filtered_libraries = [] + + for dep in ctx.attr.deps: + cc_info = dep[CcInfo] + cc_linker_inputs = cc_info.linking_context.linker_inputs + linker_input = cc_linker_inputs.to_list()[0] + + for lib in linker_input.libraries: + filtered_libraries.append(lib) + + linker_input = cc_common.create_linker_input( + owner = ctx.label, + libraries = depset(direct = filtered_libraries), + ) + + linking_context = cc_common.create_linking_context( + linker_inputs = depset(direct = [linker_input]), + ) + + return [CcInfo(linking_context = linking_context)] + +stripped_cc_info = rule( + attrs = { + "deps": attr.label_list( + allow_files = False, + providers = [CcInfo], + ), + }, + implementation = _stripped_cc_info_impl, +) diff --git a/third_party/xla/third_party/tsl/tsl/platform/build_config_root.bzl b/third_party/xla/third_party/tsl/tsl/platform/build_config_root.bzl index 151e40d4c02..fd87c70d761 100644 --- a/third_party/xla/third_party/tsl/tsl/platform/build_config_root.bzl +++ b/third_party/xla/third_party/tsl/tsl/platform/build_config_root.bzl @@ -8,6 +8,7 @@ load( _if_llvm_powerpc_available = "if_llvm_powerpc_available", _if_llvm_system_z_available = "if_llvm_system_z_available", _if_llvm_x86_available = "if_llvm_x86_available", + _if_pywrap = "if_pywrap", _if_static = "if_static", _if_static_and_not_mobile = "if_static_and_not_mobile", _tf_additional_grpc_deps_py = "tf_additional_grpc_deps_py", @@ -27,6 +28,7 @@ if_llvm_powerpc_available = _if_llvm_powerpc_available if_llvm_system_z_available = _if_llvm_system_z_available if_llvm_x86_available = _if_llvm_x86_available if_static = _if_static +if_pywrap = _if_pywrap if_static_and_not_mobile = _if_static_and_not_mobile tf_additional_grpc_deps_py = _tf_additional_grpc_deps_py tf_additional_license_deps = _tf_additional_license_deps diff --git a/third_party/xla/third_party/tsl/tsl/platform/default/BUILD b/third_party/xla/third_party/tsl/tsl/platform/default/BUILD index 444944ddecd..bcddadc908c 100644 --- a/third_party/xla/third_party/tsl/tsl/platform/default/BUILD +++ b/third_party/xla/third_party/tsl/tsl/platform/default/BUILD @@ -573,6 +573,7 @@ bzl_library( name = "build_config_root_bzl", srcs = ["build_config_root.bzl"], # copybara:uncomment parse_tests = False, + deps = ["//third_party/py/rules_pywrap:pywrap_bzl"], ) # Export source files needed for mobile builds, which do not use granular targets. diff --git a/third_party/xla/third_party/tsl/tsl/platform/default/build_config.bzl b/third_party/xla/third_party/tsl/tsl/platform/default/build_config.bzl index 4c7dcf2dcb6..b7ad02b93fc 100644 --- a/third_party/xla/third_party/tsl/tsl/platform/default/build_config.bzl +++ b/third_party/xla/third_party/tsl/tsl/platform/default/build_config.bzl @@ -8,6 +8,7 @@ load( "if_not_windows", "if_tsl_link_protobuf", ) +load("//third_party/py/rules_pywrap:pywrap.bzl", "use_pywrap_rules") load("//tsl/platform:build_config_root.bzl", "if_static") def well_known_proto_libs(): @@ -147,6 +148,8 @@ def _proto_py_outs(srcs, use_grpc_plugin = False): ret += [s[:-len(".proto")] + "_pb2_grpc.py" for s in srcs] return ret +# TODO(b/356020232): cleanup non use_pywrap_rules() parts and everythin relate +# to creation of header-only protobuf targets # Re-defined protocol buffer rule to allow building "header only" protocol # buffers, to avoid duplicate registrations. Also allows non-iterable cc_libs # containing select() statements. @@ -247,21 +250,29 @@ def cc_proto_library( }) impl_name = name + "_impl" - header_only_name = name + "_headers_only" - header_only_deps = tf_deps(protolib_deps, "_cc_headers_only") - if make_default_target_header_only: - native.alias( - name = name, - actual = header_only_name, - visibility = kwargs["visibility"], - ) - else: + if use_pywrap_rules(): native.alias( name = name, actual = impl_name, visibility = kwargs["visibility"], ) + else: + header_only_name = name + "_headers_only" + header_only_deps = tf_deps(protolib_deps, "_cc_headers_only") + + if make_default_target_header_only: + native.alias( + name = name, + actual = header_only_name, + visibility = kwargs["visibility"], + ) + else: + native.alias( + name = name, + actual = impl_name, + visibility = kwargs["visibility"], + ) native.cc_library( name = impl_name, @@ -272,14 +283,18 @@ def cc_proto_library( alwayslink = 1, **kwargs ) - native.cc_library( - name = header_only_name, - deps = [ - "@com_google_protobuf//:protobuf_headers", - ] + header_only_deps + if_tsl_link_protobuf([impl_name]), - hdrs = gen_hdrs, - **kwargs - ) + + if use_pywrap_rules(): + pass + else: + native.cc_library( + name = header_only_name, + deps = [ + "@com_google_protobuf//:protobuf_headers", + ] + header_only_deps + if_tsl_link_protobuf([impl_name]), + hdrs = gen_hdrs, + **kwargs + ) # Re-defined protocol buffer rule to allow setting service namespace. def cc_grpc_library( @@ -429,6 +444,8 @@ def py_proto_library( **kwargs ) +# TODO(b/356020232): cleanup non-use_pywrap_rules part and all logic reated to +# protobuf header-only targets after migration is done def tf_proto_library_cc( name, srcs = [], @@ -478,24 +495,36 @@ def tf_proto_library_cc( visibility = visibility, ) - native.alias( - name = cc_name + "_headers_only", - actual = cc_name, - testonly = testonly, - visibility = visibility, - ) + if use_pywrap_rules(): + pass + else: + native.alias( + name = cc_name + "_headers_only", + actual = cc_name, + testonly = testonly, + visibility = visibility, + ) + + native.cc_library( + name = cc_name, + deps = cc_deps + ["@com_google_protobuf//:protobuf_headers"] + if_tsl_link_protobuf([name + "_cc_impl"]), + testonly = testonly, + visibility = visibility, + ) - native.cc_library( - name = cc_name, - deps = cc_deps + ["@com_google_protobuf//:protobuf_headers"] + if_tsl_link_protobuf([name + "_cc_impl"]), - testonly = testonly, - visibility = visibility, - ) native.cc_library( name = cc_name + "_impl", deps = [s + "_impl" for s in cc_deps], ) + if use_pywrap_rules(): + native.cc_library( + name = cc_name, + deps = cc_deps + ["@com_google_protobuf//:protobuf_headers", cc_name + "_impl"], + testonly = testonly, + visibility = visibility, + ) + return cc_proto_library( @@ -760,6 +789,7 @@ def tf_protobuf_deps(): otherwise = [clean_dep("@com_google_protobuf//:protobuf_headers")], ) +# TODO(b/356020232): remove completely after migration is done # Link protobuf, unless the tsl_link_protobuf build flag is explicitly set to false. def tsl_protobuf_deps(): return if_tsl_link_protobuf([clean_dep("@com_google_protobuf//:protobuf")], [clean_dep("@com_google_protobuf//:protobuf_headers")]) diff --git a/third_party/xla/third_party/tsl/tsl/platform/default/build_config_root.bzl b/third_party/xla/third_party/tsl/tsl/platform/default/build_config_root.bzl index 142641b16d2..05559211a93 100644 --- a/third_party/xla/third_party/tsl/tsl/platform/default/build_config_root.bzl +++ b/third_party/xla/third_party/tsl/tsl/platform/default/build_config_root.bzl @@ -3,6 +3,7 @@ # be separate to avoid cyclic references. load("@local_config_remote_execution//:remote_execution.bzl", "gpu_test_tags") +load("//third_party/py/rules_pywrap:pywrap.bzl", "use_pywrap_rules") # RBE settings for tests that require a GPU. This is used in exec_properties of rules # that need GPU access. @@ -39,12 +40,16 @@ def tf_additional_license_deps(): def tf_additional_tpu_ops_deps(): return [] +# TODO(b/356020232): remove completely after migration is done # Include specific extra dependencies when building statically, or # another set of dependencies otherwise. If "macos" is provided, that # dependency list is used when using the framework_shared_object config # on MacOS platforms. If "macos" is not provided, the "otherwise" list is # used for all framework_shared_object platforms including MacOS. def if_static(extra_deps, otherwise = [], macos = []): + if use_pywrap_rules(): + return extra_deps + ret = { str(Label("@local_xla//xla/tsl:framework_shared_object")): otherwise, "//conditions:default": extra_deps, @@ -53,7 +58,11 @@ def if_static(extra_deps, otherwise = [], macos = []): ret[str(Label("@local_xla//xla/tsl:macos_with_framework_shared_object"))] = macos return select(ret) +# TODO(b/356020232): remove completely after migration is done def if_static_and_not_mobile(extra_deps, otherwise = []): + if use_pywrap_rules(): + return extra_deps + return select({ str(Label("@local_xla//xla/tsl:framework_shared_object")): otherwise, str(Label("@local_xla//xla/tsl:android")): otherwise, @@ -61,6 +70,10 @@ def if_static_and_not_mobile(extra_deps, otherwise = []): "//conditions:default": extra_deps, }) +# TODO(b/356020232): remove completely after migration is done +def if_pywrap(if_true = [], if_false = []): + return if_true if use_pywrap_rules() else if_false + def if_llvm_aarch32_available(then, otherwise = []): return select({ str(Label("@local_xla//xla/tsl:aarch32_or_cross")): then, diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD index 74641d71ffd..e50a634e132 100644 --- a/third_party/xla/xla/stream_executor/BUILD +++ b/third_party/xla/xla/stream_executor/BUILD @@ -31,6 +31,7 @@ bzl_library( srcs = ["build_defs.bzl"], deps = [ "@local_config_cuda//cuda:build_defs_bzl", + "@local_tsl//third_party/py/rules_pywrap:pywrap_bzl", "@local_tsl//tsl/platform:rules_cc_bzl", "@local_tsl//tsl/platform/default:cuda_build_defs_bzl", ] + stream_executor_build_defs_bzl_deps(), diff --git a/third_party/xla/xla/tsl/BUILD b/third_party/xla/xla/tsl/BUILD index e8d8f963679..ff4a95dd134 100644 --- a/third_party/xla/xla/tsl/BUILD +++ b/third_party/xla/xla/tsl/BUILD @@ -546,6 +546,7 @@ bzl_library( "@local_config_cuda//cuda:build_defs_bzl", "@local_config_rocm//rocm:build_defs_bzl", "@local_config_tensorrt//:build_defs_bzl", + "@local_tsl//third_party/py/rules_pywrap:pywrap_bzl", "@local_tsl//tsl/platform:rules_cc_bzl", ], ) diff --git a/third_party/xla/xla/tsl/tsl.bzl b/third_party/xla/xla/tsl/tsl.bzl index 1f9ce4e543f..9cff2cc413c 100644 --- a/third_party/xla/xla/tsl/tsl.bzl +++ b/third_party/xla/xla/tsl/tsl.bzl @@ -32,6 +32,11 @@ load( "@local_config_tensorrt//:build_defs.bzl", "if_tensorrt", ) +load( + "@local_tsl//third_party/py/rules_pywrap:pywrap.bzl", + "use_pywrap_rules", + _pybind_extension = "pybind_extension", +) # Internally this loads a macro, but in OSS this is a function # buildifier: disable=out-of-order-load @@ -119,7 +124,10 @@ def internal_visibility(internal_targets): return if_google(internal_targets, ["//visibility:public"]) # TODO(jakeharmon): Use this to replace if_static +# TODO(b/356020232): remove completely after migration is done def if_tsl_link_protobuf(if_true, if_false = []): + if use_pywrap_rules(): + return if_true return select({ "//conditions:default": if_true, clean_dep("//xla/tsl:tsl_protobuf_header_only"): if_false, @@ -272,6 +280,8 @@ def get_win_copts(is_external = False): return WINDOWS_COPTS + ["/DTF_COMPILE_LIBRARY"] # copybara:comment_end +# TODO(b/356020232): cleanup non-use_pywrap_rules part once migration is done +# buildozer: disable=function-docstring-args def tsl_copts( android_optimization_level_override = "-O2", is_external = False, @@ -286,6 +296,16 @@ def tsl_copts( ] if android_optimization_level_override: android_copts.append(android_optimization_level_override) + + framework_deps = [] + if use_pywrap_rules(): + pass + else: + framework_deps = select({ + clean_dep("//xla/tsl:framework_shared_object"): [], + "//conditions:default": ["-DTENSORFLOW_MONOLITHIC_BUILD"], + }) + return ( if_not_windows([ "-DEIGEN_AVOID_STL_ARRAY", @@ -314,10 +334,7 @@ def tsl_copts( if_linux_x86_64(["-msse3"]) + if_ios_x86_64(["-msse4.1"]) + if_no_default_logger(["-DNO_DEFAULT_LOGGER"]) + - select({ - clean_dep("//xla/tsl:framework_shared_object"): [], - "//conditions:default": ["-DTENSORFLOW_MONOLITHIC_BUILD"], - }) + + framework_deps + select({ clean_dep("//xla/tsl:android"): android_copts, clean_dep("//xla/tsl:emscripten"): [], @@ -497,6 +514,24 @@ def transitive_hdrs(name, deps = [], **kwargs): _transitive_hdrs(name = name + "_gather", deps = deps) native.filegroup(name = name, srcs = [":" + name + "_gather"], **kwargs) +def cc_header_only_library(name, deps = [], includes = [], extra_deps = [], compatible_with = None, **kwargs): + if use_pywrap_rules(): + cc_library( + name = name, + deps = deps + extra_deps, + compatible_with = compatible_with, + **kwargs + ) + else: + custom_op_cc_header_only_library( + name, + deps, + includes, + extra_deps, + compatible_with, + **kwargs + ) + # Create a header only library that includes all the headers exported by # the libraries in deps. # @@ -510,7 +545,7 @@ def transitive_hdrs(name, deps = [], **kwargs): # * Eigen: it's a header-only library. Add it directly to your deps. # * GRPC: add a direct dep on @com_github_grpc_grpc//:grpc++_public_hdrs. # -def cc_header_only_library(name, deps = [], includes = [], extra_deps = [], compatible_with = None, **kwargs): +def custom_op_cc_header_only_library(name, deps = [], includes = [], extra_deps = [], compatible_with = None, **kwargs): _transitive_hdrs( name = name + "_gather", deps = deps, @@ -801,3 +836,6 @@ def tsl_extra_config_settings(): def tsl_extra_config_settings_targets(): return [] + +# TODO(b/356020232): remove after migration is done +tsl_pybind_extension = _pybind_extension if use_pywrap_rules() else tsl_pybind_extension_opensource diff --git a/third_party/xla/xla/tsl/tsl.default.bzl b/third_party/xla/xla/tsl/tsl.default.bzl index 2a6e4b3a0e2..ffa4b8f9cb1 100644 --- a/third_party/xla/xla/tsl/tsl.default.bzl +++ b/third_party/xla/xla/tsl/tsl.default.bzl @@ -11,7 +11,7 @@ load( _tsl_extra_config_settings_targets = "tsl_extra_config_settings_targets", _tsl_google_bzl_deps = "tsl_google_bzl_deps", _tsl_grpc_cc_dependencies = "tsl_grpc_cc_dependencies", - _tsl_pybind_extension = "tsl_pybind_extension_opensource", + _tsl_pybind_extension = "tsl_pybind_extension", ) get_compatible_with_portable = _get_compatible_with_portable