forked from google-deepmind/deepmind-research
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathBUILD.bazel
125 lines (111 loc) · 3.38 KB
/
BUILD.bazel
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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
"""DM21 functionals."""
load("@org_tensorflow//tensorflow/python/tools:tools.bzl", "saved_model_compile_aot")
load("@rules_python//python:defs.bzl", "py_binary", "py_library")
load("@external_py_deps//:requirements.bzl", "requirement")
licenses(["notice"])
py_library(
name = "compute_hfx_density",
srcs = ["density_functional_approximation_dm21/compute_hfx_density.py"],
srcs_version = "PY3",
deps = [
requirement("numpy"),
requirement("pyscf"),
],
)
py_test(
name = "compute_hfx_density_test",
srcs = ["density_functional_approximation_dm21/compute_hfx_density_test.py"],
python_version = "PY3",
srcs_version = "PY3",
deps = [
":compute_hfx_density",
requirement("attrs"),
requirement("numpy"),
requirement("scipy"),
requirement("pyscf"),
"@io_abseil_py//absl/testing:absltest",
"@io_abseil_py//absl/testing:parameterized",
],
)
filegroup(
name = "dm21_checkpoints",
srcs = glob(["density_functional_approximation_dm21/checkpoints/**"]),
)
py_library(
name = "neural_numint",
srcs = ["density_functional_approximation_dm21/neural_numint.py"],
data = [":dm21_checkpoints"],
srcs_version = "PY3",
deps = [
":compute_hfx_density",
requirement("attrs"),
requirement("keras"),
requirement("numpy"),
requirement("pyscf"),
"@org_tensorflow//tensorflow:tensorflow_py",
requirement("tensorflow-hub"),
],
)
py_test(
name = "neural_numint_test",
srcs = ["density_functional_approximation_dm21/neural_numint_test.py"],
python_version = "PY3",
srcs_version = "PY3",
deps = [
":neural_numint",
requirement("attrs"),
requirement("pyscf"),
"@io_abseil_py//absl/testing:parameterized",
"@org_tensorflow//tensorflow:tensorflow_py",
],
)
py_binary(
name = "export_saved_model",
srcs = ["density_functional_approximation_dm21/export_saved_model.py"],
data = [":dm21_checkpoints"],
python_version = "PY3",
srcs_version = "PY3",
deps = [
":neural_numint",
],
)
EXPORTED_SAVED_MODEL_OBJECTS = [
"DM21/saved_model.pb",
"DM21/variables/variables.index",
"DM21/variables/variables.data-00000-of-00001",
]
genrule(
name = "create_model_for_aot_compile",
outs = EXPORTED_SAVED_MODEL_OBJECTS,
# The functional used in dm21_aot_compiled_example can be changed by
# setting the --functional flag to the desired functional.
cmd = "$(location :export_saved_model) --functional DM21 --batch_size 1000 --out_dir $(@D)/DM21",
tools = [":export_saved_model"],
)
filegroup(
name = "dm21_exported_model",
srcs = EXPORTED_SAVED_MODEL_OBJECTS,
)
saved_model_compile_aot(
name = "aot_compiled_dm21",
cpp_class = "dm21::functional",
directory = ":DM21",
filegroups = [":dm21_exported_model"],
force_without_xla_support_flag = False,
multithreading = False,
signature_def = "default",
tag_set = "''",
target_triple = "x86_64-pc-linux",
)
cc_library(
name = "dm21_aot_compiled_example",
srcs = ["cc/dm21_aot_compiled_example.cc"],
hdrs = ["cc/dm21_aot_compiled_example.h"],
deps = [":aot_compiled_dm21"],
)
cc_binary(
name = "run_dm21_aot_compiled_example",
srcs = ["cc/run_dm21_aot_compiled_example.cc"],
copts = ["-DXLA_AVAILABLE"],
deps = [":dm21_aot_compiled_example"],
)