xref: /aosp_15_r20/external/tensorflow/tensorflow/c/experimental/gradients/BUILD (revision b6fb3261f9314811a0f4371741dbb8839866f948) 
1load("//tensorflow:tensorflow.bzl", "filegroup")
2load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
3
4# buildifier: disable=same-origin-load
5load(
6    "//tensorflow:tensorflow.bzl",
7    "if_libtpu",
8    "tf_cuda_cc_test",
9)
10load(
11    "//tensorflow/core/platform:build_config.bzl",
12    "tf_kernel_tests_linkstatic",
13)
14load(
15    "//tensorflow/core/platform:build_config_root.bzl",
16    "tf_cuda_tests_tags",
17)
18
19# Library of gradient functions.
20package(
21    licenses = ["notice"],
22)
23
24cc_library(
25    name = "array_grad",
26    srcs = ["array_grad.cc"],
27    hdrs = [
28        "array_grad.h",
29    ],
30    visibility = [
31        "//tensorflow:internal",
32    ],
33    deps = [
34        "//tensorflow/c/eager:abstract_context",
35        "//tensorflow/c/eager:gradients_internal",
36    ],
37)
38
39cc_library(
40    name = "math_grad",
41    srcs = ["math_grad.cc"],
42    hdrs = [
43        "math_grad.h",
44    ],
45    visibility = [
46        "//tensorflow:internal",
47    ],
48    deps = [
49        "//tensorflow/c/eager:abstract_tensor_handle",
50        "//tensorflow/c/eager:gradients_internal",
51        "//tensorflow/c/experimental/ops:array_ops",
52        "//tensorflow/c/experimental/ops:math_ops",
53        "//tensorflow/c/experimental/ops:nn_ops",
54    ],
55)
56
57cc_library(
58    name = "nn_grad",
59    srcs = ["nn_grad.cc"],
60    hdrs = [
61        "nn_grad.h",
62    ],
63    visibility = [
64        "//tensorflow:internal",
65    ],
66    deps = [
67        "//tensorflow/c/eager:abstract_tensor_handle",
68        "//tensorflow/c/eager:gradients_internal",
69        "//tensorflow/c/eager:immediate_execution_context",
70        "//tensorflow/c/eager:immediate_execution_tensor_handle",
71        "//tensorflow/c/experimental/ops:array_ops",
72        "//tensorflow/c/experimental/ops:math_ops",
73        "//tensorflow/c/experimental/ops:nn_ops",
74        "//tensorflow/core/lib/llvm_rtti",
75        "//tensorflow/core/platform:errors",
76        "@com_google_absl//absl/types:span",
77    ],
78)
79
80cc_library(
81    name = "not_differentiable",
82    srcs = ["not_differentiable.cc"],
83    hdrs = [
84        "not_differentiable.h",
85    ],
86    visibility = [
87        "//tensorflow:internal",
88    ],
89    deps = [
90        "//tensorflow/c/eager:abstract_context",
91        "//tensorflow/c/eager:gradients_internal",
92    ],
93)
94
95cc_library(
96    name = "gradients",
97    hdrs = [
98        "array_grad.h",
99        "math_grad.h",
100        "nn_grad.h",
101        "not_differentiable.h",
102    ],
103    visibility = [
104        "//tensorflow:internal",
105    ],
106    deps = [
107        ":array_grad",
108        ":math_grad",
109        ":nn_grad",
110        ":not_differentiable",
111        "//tensorflow/c/eager:abstract_context",
112        "//tensorflow/c/eager:gradients_internal",
113    ],
114)
115
116tf_cuda_cc_test(
117    name = "custom_gradient_test",
118    size = "small",
119    srcs = [
120        "custom_gradient_test.cc",
121    ],
122    args = ["--heap_check="],  # TODO(b/174752220): Remove
123    linkstatic = tf_kernel_tests_linkstatic(),
124    tags = tf_cuda_tests_tags(),
125    deps = [
126        "//tensorflow/c:tf_status_helper",
127        "//tensorflow/c/eager:abstract_context",
128        "//tensorflow/c/eager:c_api",
129        "//tensorflow/c/eager:c_api_experimental",
130        "//tensorflow/c/eager:c_api_unified_internal",
131        "//tensorflow/c/eager:gradients_internal",
132        "//tensorflow/c/eager:unified_api_testutil",
133        "//tensorflow/c/experimental/ops",
134        "//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration",
135        "//tensorflow/core:test",
136        "//tensorflow/core:test_main",
137        "//tensorflow/core/platform:errors",
138    ],
139)
140
141filegroup(
142    name = "pywrap_required_hdrs",
143    srcs = [
144        "array_grad.h",
145        "math_grad.h",
146        "nn_grad.h",
147        "not_differentiable.h",
148    ],
149    visibility = [
150        "//tensorflow/core:__pkg__",
151        "//tensorflow/python:__pkg__",
152    ],
153)
154
155cc_library(
156    name = "grad_test_helper",
157    testonly = True,
158    srcs = ["grad_test_helper.cc"],
159    hdrs = ["grad_test_helper.h"],
160    visibility = [
161        "//tensorflow:internal",
162    ],
163    deps = [
164        "//tensorflow/c/eager:gradient_checker",
165        "//tensorflow/c/eager:gradients_internal",
166        "//tensorflow/c/eager:unified_api_testutil",
167        "//tensorflow/c/experimental/gradients/tape:tape_context",
168        "//tensorflow/core:test",
169        "//tensorflow/core:test_main",
170    ],
171)
172
173tf_cuda_cc_test(
174    name = "nn_grad_test",
175    size = "small",
176    srcs = [
177        "nn_grad_test.cc",
178    ],
179    args = ["--heap_check="],  # TODO(b/174752220): Remove
180    linkstatic = tf_kernel_tests_linkstatic(),
181    tags = tf_cuda_tests_tags() + ["no_cuda_asan"],  # b/173654156,
182    deps = [
183        ":grad_test_helper",
184        ":nn_grad",
185        "//tensorflow/c:tf_status_helper",
186        "//tensorflow/c/eager:c_api_test_util",
187        "//tensorflow/c/eager:c_api_unified_internal",
188        "//tensorflow/c/eager:unified_api_testutil",
189        "//tensorflow/c/experimental/gradients/tape:tape_context",
190        "//tensorflow/c/experimental/ops:nn_ops",
191        "//tensorflow/core/platform:tensor_float_32_utils",
192        "//tensorflow/core:test",
193        "//tensorflow/core:test_main",
194    ] + if_libtpu(
195        if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"],
196        if_true = [],
197    ),
198)
199
200tf_cuda_cc_test(
201    name = "math_grad_test",
202    size = "small",
203    srcs = [
204        "math_grad_test.cc",
205    ],
206    args = ["--heap_check="],  # TODO(b/174752220): Remove
207    linkstatic = tf_kernel_tests_linkstatic(),
208    tags = tf_cuda_tests_tags() + ["no_cuda_asan"],  # b/173654156,
209    deps = [
210        ":grad_test_helper",
211        ":math_grad",
212        "//tensorflow/c:tf_status_helper",
213        "//tensorflow/c/eager:c_api_test_util",
214        "//tensorflow/c/eager:c_api_unified_internal",
215        "//tensorflow/c/eager:unified_api_testutil",
216        "//tensorflow/c/experimental/gradients/tape:tape_context",
217        "//tensorflow/c/experimental/ops:math_ops",
218        "//tensorflow/core/platform:tensor_float_32_utils",
219        "//tensorflow/core:test",
220        "//tensorflow/core:test_main",
221    ] + if_libtpu(
222        if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"],
223        if_true = [],
224    ),
225)
226
227tf_cuda_cc_test(
228    name = "array_grad_test",
229    size = "small",
230    srcs = [
231        "array_grad_test.cc",
232    ],
233    args = ["--heap_check="],  # TODO(b/174752220): Remove
234    linkstatic = tf_kernel_tests_linkstatic(),
235    tags = tf_cuda_tests_tags() + ["no_cuda_asan"],  # b/173654156,
236    deps = [
237        ":grad_test_helper",
238        ":array_grad",
239        "//tensorflow/c:tf_status_helper",
240        "//tensorflow/c/eager:c_api_test_util",
241        "//tensorflow/c/experimental/gradients/tape:tape_context",
242        "//tensorflow/c/experimental/ops:array_ops",
243        "//tensorflow/core/platform:tensor_float_32_utils",
244        "//tensorflow/core:test",
245        "//tensorflow/core:test_main",
246        "//tensorflow/c/eager:c_api_unified_internal",
247        "//tensorflow/c/eager:unified_api_testutil",
248    ] + if_libtpu(
249        if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"],
250        if_true = [],
251    ),
252)
253