forked from tensorflow/tensorflow
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathgen_op.bzl
More file actions
65 lines (59 loc) · 2.09 KB
/
gen_op.bzl
File metadata and controls
65 lines (59 loc) · 2.09 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
"""For eager-mode Python."""
load("//tensorflow:tensorflow.bzl",
"clean_dep",
"tf_binary_additional_srcs",
"tf_copts",
"tf_cc_binary")
def tfe_gen_op_wrapper_py(name,
out=None,
visibility=None,
deps=[],
generated_target_name=None,
# ApiDefs will be loaded in the order specified in this list.
api_def_srcs=[]):
"""Generate an eager-mode Python op wrapper for an op library."""
# Construct a cc_binary containing the specified ops.
tool_name = "gen_" + name + "_py_wrappers_cc"
if not deps:
deps = [str(Label("//tensorflow/core:" + name + "_op_lib"))]
tf_cc_binary(
name=tool_name,
linkopts=["-lm"],
copts=tf_copts(),
linkstatic=1,
deps=([
clean_dep("//tensorflow/python/eager:python_eager_op_gen_main")
] + deps),
visibility=[clean_dep("//visibility:public")],)
# Invoke the previous cc_binary to generate a python file.
if not out:
out = "gen_" + name + ".py"
if not api_def_srcs:
api_def_args_str = ","
else:
api_def_args = []
for api_def_src in api_def_srcs:
# Add directory of the first ApiDef source to args.
# We are assuming all ApiDefs in a single api_def_src are in the
# same directory.
api_def_args.append(
"$$(dirname $$(echo $(locations " + api_def_src +
") | cut -d\" \" -f1))")
api_def_args_str = ",".join(api_def_args)
native.genrule(
name=name + "_pygenrule",
outs=[out],
srcs=api_def_srcs,
tools=[tool_name] + tf_binary_additional_srcs(),
cmd=("$(location " + tool_name + ") " + api_def_args_str + " > $@"))
# Make a py_library out of the generated python file.
if not generated_target_name:
generated_target_name = name
native.py_library(
name=generated_target_name,
srcs=[out],
srcs_version="PY2AND3",
visibility=visibility,
deps=[
clean_dep("//tensorflow/python/eager:framework_for_generated_wrappers"),
],)