Skip to content

Commit ea867d6

Browse files
ddunltensorflower-gardener
ddunl
authored and
tensorflower-gardener
committed
Move tsl/python to xla/tsl/python
PiperOrigin-RevId: 620320903
1 parent f606d83 commit ea867d6

35 files changed

+83
-75
lines changed

‎tensorflow/lite/python/interpreter_wrapper/numpy.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,8 @@ limitations under the License.
4040
// translation unit boundaries.
4141
//
4242
// For more info see https://sourceforge.net/p/numpy/mailman/message/5700519
43-
// See also tensorflow/tsl/python/lib/core/numpy.h for a similar approach.
43+
// See also tensorflow/compiler/xla/tsl/python/lib/core/numpy.h for a similar
44+
// approach.
4445
#define PY_ARRAY_UNIQUE_SYMBOL _tflite_numpy_api
4546
#ifndef TFLITE_IMPORT_NUMPY
4647
#define NO_IMPORT_ARRAY

‎tensorflow/python/BUILD

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -834,9 +834,9 @@ pywrap_tensorflow_macro(
834834
"@com_google_absl//absl/types:span",
835835
"@local_tsl//tsl/profiler/rpc:profiler_server_impl",
836836
"@local_tsl//tsl/profiler/rpc/client:profiler_client_impl",
837-
"@local_tsl//tsl/python/lib/core:numpy",
838837
"@local_xla//xla/backends/profiler/cpu:python_tracer",
839838
"@local_xla//xla/stream_executor:stream_executor_impl",
839+
"@local_xla//xla/tsl/python/lib/core:numpy",
840840
] + select({
841841
"//tensorflow/compiler/mlir/python:disable_mlir_config": [],
842842
"//conditions:default": [
@@ -951,9 +951,9 @@ filegroup(
951951
"@local_tsl//tsl/profiler/lib:profiler_session_impl", # profiler
952952
"@local_tsl//tsl/profiler/rpc:profiler_server_impl", # profiler
953953
"@local_tsl//tsl/profiler/rpc/client:profiler_client_impl",
954-
"@local_tsl//tsl/python/lib/core:ml_dtypes_lib", # bfloat16, float8_e4m3fn, float8_e5m2
955-
"@local_tsl//tsl/python/lib/core:numpy", # checkpoint_reader
956954
"@local_xla//xla/stream_executor", # stat_summarizer
955+
"@local_xla//xla/tsl/python/lib/core:ml_dtypes_lib", # bfloat16, float8_e4m3fn, float8_e5m2
956+
"@local_xla//xla/tsl/python/lib/core:numpy", # checkpoint_reader
957957
] + select({
958958
"//tensorflow/compiler/mlir/python:disable_mlir_config": [],
959959
"//conditions:default": [
@@ -1151,7 +1151,7 @@ cc_library(
11511151
name = "unified_api_pywrap_required_headers",
11521152
textual_hdrs = [
11531153
"//tensorflow/python/lib/core:basic_hdrs",
1154-
"@local_tsl//tsl/python/lib/core:basic_hdrs",
1154+
"@local_xla//xla/tsl/python/lib/core:basic_hdrs",
11551155
"//tensorflow/c:headers",
11561156
"//tensorflow/c:safe_ptr_hdr",
11571157
"//tensorflow/c/eager:headers",
@@ -1205,7 +1205,7 @@ tf_python_pybind_extension(
12051205
"//tensorflow/python/util:util_hdr",
12061206
"@local_tsl//tsl/distributed_runtime:pywrap_required_hdrs",
12071207
"@local_tsl//tsl/distributed_runtime/coordination:pywrap_required_hdrs",
1208-
"@local_tsl//tsl/python/lib/core:numpy_hdr",
1208+
"@local_xla//xla/tsl/python/lib/core:numpy_hdr",
12091209
],
12101210
dynamic_deps = [":_pywrap_tensorflow_internal.so"] + select({
12111211
"//tensorflow:macos": ["//tensorflow:libtensorflow_framework.%s.dylib" % VERSION],

‎tensorflow/python/client/BUILD

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ tf_python_pybind_extension(
4545
"//tensorflow/python/lib/core:safe_pyobject_ptr_required_hdrs",
4646
"@local_tsl//tsl/distributed_runtime:pywrap_required_hdrs",
4747
"@local_tsl//tsl/distributed_runtime/coordination:pywrap_required_hdrs",
48-
"@local_tsl//tsl/python/lib/core:numpy_hdr",
48+
"@local_xla//xla/tsl/python/lib/core:numpy_hdr",
4949
],
5050
enable_stub_generation = True,
5151
pytype_srcs = [
@@ -68,7 +68,7 @@ tf_python_pybind_extension(
6868
"//third_party/python_runtime:headers",
6969
"@com_google_absl//absl/types:optional",
7070
"@eigen_archive//:eigen3",
71-
"@local_tsl//tsl/python/lib/core:numpy",
71+
"@local_xla//xla/tsl/python/lib/core:numpy",
7272
"@pybind11",
7373
"@pybind11_abseil//pybind11_abseil:absl_casters",
7474
"@pybind11_protobuf//pybind11_protobuf:native_proto_caster",
@@ -259,7 +259,7 @@ tf_cuda_library(
259259
"//tensorflow/python/lib/core:safe_pyobject_ptr",
260260
"@com_google_absl//absl/status",
261261
"@com_google_absl//absl/strings",
262-
"@local_tsl//tsl/python/lib/core:numpy",
262+
"@local_xla//xla/tsl/python/lib/core:numpy",
263263
],
264264
alwayslink = 1,
265265
)

‎tensorflow/python/client/tf_session_helper.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ limitations under the License.
1818

1919
// Must be included first
2020
// clang-format off
21-
#include "tsl/python/lib/core/numpy.h" //NOLINT
21+
#include "xla/tsl/python/lib/core/numpy.h" //NOLINT
2222
// clang-format on
2323

2424
#include "tensorflow/c/c_api.h"

‎tensorflow/python/client/tf_session_wrapper.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ limitations under the License.
4040
#include "tensorflow/c/safe_ptr.h"
4141
#include "tensorflow/c/tf_buffer.h"
4242
#include "tensorflow/c/tf_datatype.h"
43+
#include "xla/tsl/python/lib/core/numpy.h"
4344
#include "tensorflow/core/distributed_runtime/server_lib.h"
4445
#include "tensorflow/core/framework/full_type.pb.h"
4546
#include "tensorflow/core/framework/versions.pb.h"
@@ -50,7 +51,6 @@ limitations under the License.
5051
#include "tensorflow/python/lib/core/pybind11_status.h"
5152
#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
5253
#include "tsl/platform/mutex.h"
53-
#include "tsl/python/lib/core/numpy.h"
5454

5555
namespace pybind11 {
5656
namespace detail {

‎tensorflow/python/eager/BUILD

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ cc_library(
8080
"@com_google_absl//absl/types:variant",
8181
"@local_tsl//tsl/platform:status",
8282
"@local_tsl//tsl/profiler/lib:traceme",
83-
"@local_tsl//tsl/python/lib/core:numpy",
83+
"@local_xla//xla/tsl/python/lib/core:numpy",
8484
"@pybind11",
8585
],
8686
)

‎tensorflow/python/eager/pywrap_tensor.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ limitations under the License.
1414
==============================================================================*/
1515
// Must be included first
1616
// clang-format off
17-
#include "tsl/python/lib/core/numpy.h" //NOLINT
17+
#include "xla/tsl/python/lib/core/numpy.h" //NOLINT
1818
// clang-format on
1919

2020
#include "tensorflow/python/eager/pywrap_tensor.h"

‎tensorflow/python/eager/pywrap_tensor.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ limitations under the License.
1717

1818
// Must be included first
1919
// clang-format off
20-
#include "tsl/python/lib/core/numpy.h" //NOLINT
20+
#include "xla/tsl/python/lib/core/numpy.h" //NOLINT
2121
// clang-format on
2222

2323
#include "tensorflow/c/eager/c_api.h"

‎tensorflow/python/framework/BUILD

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -865,7 +865,7 @@ tf_python_pybind_extension(
865865
"//tensorflow/python/lib/core:safe_pyobject_ptr_required_hdrs",
866866
"@local_tsl//tsl/distributed_runtime:pywrap_required_hdrs",
867867
"@local_tsl//tsl/distributed_runtime/coordination:pywrap_required_hdrs",
868-
"@local_tsl//tsl/python/lib/core:numpy_hdr",
868+
"@local_xla//xla/tsl/python/lib/core:numpy_hdr",
869869
],
870870
enable_stub_generation = True,
871871
pytype_srcs = [
@@ -965,7 +965,7 @@ tf_python_pybind_extension(
965965
"//tensorflow/python/lib/core:safe_pyobject_ptr_required_hdrs",
966966
"@local_tsl//tsl/distributed_runtime:pywrap_required_hdrs",
967967
"@local_tsl//tsl/distributed_runtime/coordination:pywrap_required_hdrs",
968-
"@local_tsl//tsl/python/lib/core:numpy_hdr",
968+
"@local_xla//xla/tsl/python/lib/core:numpy_hdr",
969969
],
970970
enable_stub_generation = True,
971971
pytype_srcs = [
@@ -1109,7 +1109,7 @@ tf_python_pybind_extension(
11091109
"//tensorflow/python/lib/core:safe_pyobject_ptr_required_hdrs",
11101110
"@local_tsl//tsl/distributed_runtime:pywrap_required_hdrs",
11111111
"@local_tsl//tsl/distributed_runtime/coordination:pywrap_required_hdrs",
1112-
"@local_tsl//tsl/python/lib/core:numpy_hdr",
1112+
"@local_xla//xla/tsl/python/lib/core:numpy_hdr",
11131113
],
11141114
enable_stub_generation = True,
11151115
pytype_srcs = [

‎tensorflow/python/lib/core/BUILD

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,8 @@ cc_library(
3939
"//tensorflow/c:tf_datatype_hdrs",
4040
"//tensorflow/core:lib",
4141
"//tensorflow/core:protos_all_cc",
42-
"@local_tsl//tsl/python/lib/core:ml_dtypes_lib",
43-
"@local_tsl//tsl/python/lib/core:numpy",
42+
"@local_xla//xla/tsl/python/lib/core:ml_dtypes_lib",
43+
"@local_xla//xla/tsl/python/lib/core:numpy",
4444
],
4545
)
4646

@@ -171,7 +171,7 @@ cc_library(
171171
"//tensorflow/python/eager:pywrap_tfe_lib",
172172
"//third_party/py/numpy:headers",
173173
"//third_party/python_runtime:headers",
174-
"@local_tsl//tsl/python/lib/core:numpy",
174+
"@local_xla//xla/tsl/python/lib/core:numpy",
175175
],
176176
alwayslink = 1,
177177
)
@@ -213,7 +213,7 @@ cc_library(
213213
"//tensorflow/c:headers",
214214
"//tensorflow/c:safe_ptr_hdr",
215215
"//tensorflow/c/eager:headers",
216-
"@local_tsl//tsl/python/lib/core:numpy_hdr",
216+
"@local_xla//xla/tsl/python/lib/core:numpy_hdr",
217217
],
218218
features = [
219219
"-parse_headers",
@@ -228,7 +228,7 @@ cc_library(
228228
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
229229
"//third_party/py/numpy:headers",
230230
"//third_party/python_runtime:headers",
231-
"@local_tsl//tsl/python/lib/core:numpy",
231+
"@local_xla//xla/tsl/python/lib/core:numpy",
232232
],
233233
)
234234

@@ -249,8 +249,8 @@ cc_library(
249249
"//tensorflow/core:framework",
250250
"//tensorflow/core:lib",
251251
"//tensorflow/python/lib/core:safe_pyobject_ptr",
252-
"@local_tsl//tsl/python/lib/core:ml_dtypes_lib",
253-
"@local_tsl//tsl/python/lib/core:numpy",
252+
"@local_xla//xla/tsl/python/lib/core:ml_dtypes_lib",
253+
"@local_xla//xla/tsl/python/lib/core:numpy",
254254
],
255255
)
256256

@@ -273,7 +273,7 @@ cc_library(
273273
"//tensorflow/core:framework",
274274
"//tensorflow/core:lib",
275275
"//third_party/python_runtime:headers", # build_cleaner: keep; DNR: b/35864863
276-
"@local_tsl//tsl/python/lib/core:numpy",
276+
"@local_xla//xla/tsl/python/lib/core:numpy",
277277
],
278278
)
279279

‎tensorflow/python/lib/core/ndarray_tensor.cc

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ limitations under the License.
1515
// Must be included first
1616
// clang-format off
1717
#include "tensorflow/c/tf_datatype.h"
18-
#include "tsl/python/lib/core/numpy.h" //NOLINT
18+
#include "xla/tsl/python/lib/core/numpy.h" //NOLINT
1919
// clang-format on
2020

2121
#include "tensorflow/python/lib/core/ndarray_tensor.h"
@@ -25,14 +25,14 @@ limitations under the License.
2525

2626
#include "tensorflow/c/eager/tfe_context_internal.h"
2727
#include "tensorflow/c/tf_tensor_internal.h"
28+
#include "xla/tsl/python/lib/core/ml_dtypes.h"
2829
#include "tensorflow/core/lib/core/coding.h"
2930
#include "tensorflow/core/lib/core/errors.h"
3031
#include "tensorflow/core/lib/gtl/inlined_vector.h"
3132
#include "tensorflow/core/platform/types.h"
3233
#include "tensorflow/core/util/port.h"
3334
#include "tensorflow/python/lib/core/ndarray_tensor_bridge.h"
3435
#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
35-
#include "tsl/python/lib/core/ml_dtypes.h"
3636

3737
namespace tensorflow {
3838
namespace {

‎tensorflow/python/lib/core/ndarray_tensor_bridge.cc

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,18 +16,18 @@ limitations under the License.
1616
// clang-format off
1717
// Must be included first.
1818
#include "tensorflow/c/tf_datatype.h"
19-
#include "tsl/python/lib/core/numpy.h"
19+
#include "xla/tsl/python/lib/core/numpy.h"
2020
// clang-format on
2121

2222
#include "tensorflow/python/lib/core/ndarray_tensor_bridge.h"
2323

2424
#include <vector>
2525

2626
#include "tensorflow/c/c_api.h"
27+
#include "xla/tsl/python/lib/core/ml_dtypes.h"
2728
#include "tensorflow/core/lib/core/errors.h"
2829
#include "tensorflow/core/platform/mutex.h"
2930
#include "tensorflow/python/lib/core/py_util.h"
30-
#include "tsl/python/lib/core/ml_dtypes.h"
3131

3232
namespace tensorflow {
3333

‎tensorflow/python/lib/core/ndarray_tensor_bridge.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ limitations under the License.
1717

1818
// Must be included first
1919
// clang-format off
20-
#include "tsl/python/lib/core/numpy.h" //NOLINT
20+
#include "xla/tsl/python/lib/core/numpy.h" //NOLINT
2121
// clang-format on
2222

2323
#include <functional>

‎tensorflow/python/lib/core/py_func.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ limitations under the License.
1717
// Must be included first.
1818
#include "tensorflow/python/lib/core/py_func.h"
1919

20-
#include "tsl/python/lib/core/numpy.h"
20+
#include "xla/tsl/python/lib/core/numpy.h"
2121
// clang-format: on
2222

2323
#include <Python.h>

‎tensorflow/python/lib/core/py_seq_tensor.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ limitations under the License.
1414
==============================================================================*/
1515
// Must be included first
1616
// clang-format off
17-
#include "tsl/python/lib/core/numpy.h" //NOLINT
17+
#include "xla/tsl/python/lib/core/numpy.h" //NOLINT
1818
// clang-format on
1919

2020
#include "tensorflow/python/lib/core/py_seq_tensor.h"

‎tensorflow/python/tfe_wrapper.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ limitations under the License.
2020
#include "absl/status/status.h"
2121
#include "absl/strings/str_cat.h"
2222
#include "pybind11/attr.h" // from @pybind11
23-
#include "tsl/python/lib/core/numpy.h" //NOLINT
23+
#include "xla/tsl/python/lib/core/numpy.h" //NOLINT
2424
// clang-format on
2525

2626
#include "Python.h"

‎tensorflow/tools/def_file_filter/symbols_pybind.txt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ tensorflow::tfprof::SerializeToString
5959
[//tensorflow/core/grappler/graph_analyzer:graph_analyzer_tool] # graph_analyze
6060
tensorflow::grappler::graph_analyzer::GraphAnalyzerTool
6161

62-
[//external/local_tsl/tsl/python/lib/core:ml_dtypes_lib] # bfloat16, float8
62+
[//external/local_xla/xla/tsl/python/lib/core:ml_dtypes_lib] # bfloat16, float8
6363
tsl::ml_dtypes::RegisterTypes
6464
tsl::ml_dtypes::GetBfloat16Dtype
6565
tsl::ml_dtypes::GetFloat8E4m3b11fnuzDtype
@@ -307,7 +307,7 @@ tensorflow::AddWhileInputHack
307307
tensorflow::RecordMutation
308308
tensorflow::Graph::IsControlEdge
309309

310-
[//external/local_tsl/tsl/python/lib/core:numpy] # tf_session
310+
[//external/local_xla/xla/tsl/python/lib/core:numpy] # tf_session
311311
tsl::ImportNumpy
312312
_tsl_numpy_api
313313

‎tensorflow/tools/pip_package/build_pip_package.py

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -219,14 +219,21 @@ def patch_so(srcs_dir: str) -> None:
219219
srcs_dir: target directory with .so files to patch.
220220
"""
221221
to_patch = {
222-
"tensorflow/python/_pywrap_tensorflow_internal.so":
223-
"$ORIGIN/../../tensorflow/tsl/python/lib/core",
224-
("tensorflow/compiler/mlir/quantization/tensorflow/python/"
225-
"pywrap_function_lib.so"): "$ORIGIN/../../../../../python",
226-
("tensorflow/compiler/mlir/quantization/tensorflow/python/"
227-
"pywrap_quantize_model.so"): "$ORIGIN/../../../../../python",
228-
("tensorflow/compiler/mlir/quantization/tensorflow/calibrator/"
229-
"pywrap_calibration.so"): "$ORIGIN/../../../../../python",
222+
"tensorflow/python/_pywrap_tensorflow_internal.so": (
223+
"$ORIGIN/../../tensorflow/compiler/xla/tsl/python/lib/core"
224+
),
225+
(
226+
"tensorflow/compiler/mlir/quantization/tensorflow/python/"
227+
"pywrap_function_lib.so"
228+
): "$ORIGIN/../../../../../python",
229+
(
230+
"tensorflow/compiler/mlir/quantization/tensorflow/python/"
231+
"pywrap_quantize_model.so"
232+
): "$ORIGIN/../../../../../python",
233+
(
234+
"tensorflow/compiler/mlir/quantization/tensorflow/calibrator/"
235+
"pywrap_calibration.so"
236+
): "$ORIGIN/../../../../../python",
230237
}
231238
for file, path in to_patch.items():
232239
rpath = subprocess.check_output(

‎third_party/xla/third_party/tsl/tools/def_file_filter/symbols_pybind.txt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ tensorflow::tfprof::SerializeToString
5959
[//tensorflow/core/grappler/graph_analyzer:graph_analyzer_tool] # graph_analyze
6060
tensorflow::grappler::graph_analyzer::GraphAnalyzerTool
6161

62-
[//external/local_tsl/tsl/python/lib/core:ml_dtypes_lib] # bfloat16, float8
62+
[//external/local_xla/xla/tsl/python/lib/core:ml_dtypes_lib] # bfloat16, float8
6363
tsl::ml_dtypes::RegisterTypes
6464
tsl::ml_dtypes::GetBfloat16Dtype
6565
tsl::ml_dtypes::GetFloat8E4m3b11fnuzDtype
@@ -307,7 +307,7 @@ tensorflow::AddWhileInputHack
307307
tensorflow::RecordMutation
308308
tensorflow::Graph::IsControlEdge
309309

310-
[//external/local_tsl/tsl/python/lib/core:numpy] # tf_session
310+
[//external/local_xla/xla/tsl/python/lib/core:numpy] # tf_session
311311
tsl::ImportNumpy
312312
_tsl_numpy_api
313313

‎third_party/xla/third_party/tsl/tsl/platform/ml_dtypes.h

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -20,16 +20,16 @@ limitations under the License.
2020
#include "ml_dtypes/include/int4.h" // from @ml_dtypes
2121

2222
namespace tsl {
23-
using float8_e4m3fn = ml_dtypes::float8_e4m3fn;
24-
using float8_e4m3fnuz = ml_dtypes::float8_e4m3fnuz;
25-
using float8_e4m3b11fnuz = ml_dtypes::float8_e4m3b11fnuz;
23+
using float8_e4m3fn = ::ml_dtypes::float8_e4m3fn;
24+
using float8_e4m3fnuz = ::ml_dtypes::float8_e4m3fnuz;
25+
using float8_e4m3b11fnuz = ::ml_dtypes::float8_e4m3b11fnuz;
2626
using float8_e4m3b11 = float8_e4m3b11fnuz; // Deprecated: old name for
2727
// backward-compatibility only.
28-
using float8_e5m2 = ml_dtypes::float8_e5m2;
29-
using float8_e5m2fnuz = ml_dtypes::float8_e5m2fnuz;
28+
using float8_e5m2 = ::ml_dtypes::float8_e5m2;
29+
using float8_e5m2fnuz = ::ml_dtypes::float8_e5m2fnuz;
3030

31-
using int4 = ml_dtypes::int4;
32-
using uint4 = ml_dtypes::uint4;
31+
using int4 = ::ml_dtypes::int4;
32+
using uint4 = ::ml_dtypes::uint4;
3333
} // namespace tsl
3434

3535
#endif // TENSORFLOW_TSL_PLATFORM_ML_DTYPES_H_

0 commit comments

Comments
 (0)