xref: /aosp_15_r20/external/tensorflow/tensorflow/python/compiler/xla/BUILD (revision b6fb3261f9314811a0f4371741dbb8839866f948) 
1load("//tensorflow:tensorflow.bzl", "cuda_py_test")
2
3package(
4    default_visibility = ["//visibility:public"],
5    licenses = ["notice"],
6)
7
8py_library(
9    name = "compiler_py",
10    srcs = [
11        "__init__.py",
12        "jit.py",
13    ],
14    srcs_version = "PY3",
15    deps = [
16        ":xla",
17        "//tensorflow/core:protos_all_py",
18        "//tensorflow/python:framework_for_generated_wrappers",
19        "//tensorflow/python/eager:context",
20    ],
21)
22
23cuda_py_test(
24    name = "jit_test",
25    size = "small",
26    srcs = ["jit_test.py"],
27    tags = [
28        "no_windows",  # TODO(b/171385770)
29    ],
30    xla_enabled = True,
31    deps = [
32        ":compiler_py",
33        "//tensorflow/python:array_ops",
34        "//tensorflow/python:client",
35        "//tensorflow/python:client_testlib",
36        "//tensorflow/python:framework",
37        "//tensorflow/python:framework_for_generated_wrappers",
38        "//tensorflow/python:framework_test_lib",
39        "//tensorflow/python:gradients",
40        "//tensorflow/python:init_ops",
41        "//tensorflow/python:math_ops",
42        "//tensorflow/python:platform_test",
43        "//tensorflow/python:random_ops",
44        "//tensorflow/python:variable_scope",
45        "//tensorflow/python:variables",
46        "//third_party/py/numpy",
47    ],
48)
49
50py_library(
51    name = "xla",
52    srcs = ["xla.py"],
53    srcs_version = "PY3",
54    deps = [
55        "//tensorflow/compiler/jit:xla_ops_py",
56        "//tensorflow/compiler/jit/ops:xla_ops_grad",
57        # Do not remove: required to run xla ops on Cloud.
58        "//tensorflow/compiler/tf2xla/python:xla",
59        "//tensorflow/python:array_ops",
60        "//tensorflow/python:control_flow_ops",
61        "//tensorflow/python:framework_ops",
62        "//tensorflow/python:platform",
63        "//tensorflow/python:util",
64        "//tensorflow/python:variable_scope",
65        "//tensorflow/python/distribute:summary_op_util",
66        "//tensorflow/python/eager:context",
67        "//tensorflow/python/eager:def_function",
68    ],
69)
70
71cuda_py_test(
72    name = "xla_test",
73    srcs = ["xla_test.py"],
74    tags = [
75        "no_mac",
76        "no_windows",
77    ],
78    xla_enabled = True,
79    deps = [
80        ":xla",
81        "//tensorflow/compiler/tests:xla_test",
82        "//tensorflow/python:client_testlib",
83        "//tensorflow/python:constant_op",
84        "//tensorflow/python:control_flow_ops",
85        "//tensorflow/python:control_flow_util",
86        "//tensorflow/python:math_ops",
87        "//tensorflow/python:platform",
88        "//tensorflow/python:state_ops",
89        "//tensorflow/python:variable_scope",
90        "//tensorflow/python/estimator:estimator_py",
91        "//tensorflow/python/summary",
92        "//tensorflow/python/tpu:tpu_lib",
93        "@absl_py//absl/testing:parameterized",
94    ],
95)
96
97cuda_py_test(
98    name = "jit_compile_test",
99    srcs = ["jit_compile_test.py"],
100    python_version = "PY3",
101    tags = [
102        "no_mac",
103        "no_windows",
104    ],
105    xla_enabled = True,
106    deps = [
107        "//tensorflow/python:client_testlib",
108        "//tensorflow/python:constant_op",
109        "//tensorflow/python:framework_ops",
110        "//tensorflow/python:resource_variable_ops",
111    ],
112)
113