61 lines
1.2 KiB
Python
61 lines
1.2 KiB
Python
licenses(["notice"]) # Apache 2.0
|
|
|
|
package(default_visibility = ["//visibility:public"])
|
|
|
|
py_library(
|
|
name = "variable_mgr_util",
|
|
srcs = [
|
|
"variable_mgr_util.py",
|
|
],
|
|
deps = [
|
|
"//third_party/py/tensorflow",
|
|
"//third_party/tensorflow/python:data_flow_ops",
|
|
"//third_party/tensorflow/python:framework_ops",
|
|
],
|
|
)
|
|
|
|
py_library(
|
|
name = "allreduce",
|
|
srcs = [
|
|
"allreduce.py",
|
|
],
|
|
deps = [
|
|
"//third_party/py/tensorflow",
|
|
"//third_party/tensorflow/contrib/all_reduce",
|
|
"//third_party/tensorflow/python:collective_ops",
|
|
"//third_party/tensorflow/python:nccl_ops",
|
|
],
|
|
)
|
|
|
|
py_library(
|
|
name = "batch_allreduce",
|
|
srcs = [
|
|
"batch_allreduce.py",
|
|
],
|
|
deps = [
|
|
":allreduce",
|
|
":constants",
|
|
"//third_party/py/tensorflow",
|
|
],
|
|
)
|
|
|
|
py_library(
|
|
name = "variable_mgr",
|
|
srcs = [
|
|
"variable_mgr.py",
|
|
],
|
|
deps = [
|
|
":allreduce",
|
|
":batch_allreduce",
|
|
":variable_mgr_util",
|
|
"//third_party/py/tensorflow",
|
|
],
|
|
)
|
|
|
|
py_library(
|
|
name = "constants",
|
|
srcs = [
|
|
"constants.py",
|
|
],
|
|
)
|