experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Rename TPUReplicate op to _TPUReplicate so it won't be accidentally serialized.

Browse Source 
Renaming this op is fine because it's only a transient op occurring in GraphOptimizationPass and it's never serialized in GraphDef.

PiperOrigin-RevId: 249517076
This commit is contained in:
Tong Shen 2019-05-22 14:11:37 -07:00 committed by TensorFlower Gardener
parent 71bc1d5198
commit 9dcf320025
7 changed files with 3 additions and 545 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/compiler/tf2xla/functionalize_control_flow.cc
View File

@ -247,8 +247,8 @@ Status FunctionalizeControlFlowPass::Run(
// multiple times, and we want to avoid functionalize it again.
static std::map<string, string>* kNodeTypeToFunctionAttrMapping =
new std::map<string, string>{
// TPUReplicate ops are generated by EncapsulateTPUComputationsPass.
{"TPUReplicate", "computation"},
// _TPUReplicate ops are generated by EncapsulateTPUComputationsPass.
{"_TPUReplicate", "computation"},
// XlaLaunch ops are generated by EncapsulateXlaComputationsPass.
{"XlaLaunch", "function"},
};

100
tensorflow/core/api_def/base_api/api_def_TPUReplicate.pbtxt
View File

@ -1,100 +0,0 @@
op {
graph_op_name: "TPUReplicate"
visibility: HIDDEN
in_arg {
name: "inputs"
description: <<END
the inputs to 'computation', flattened, in replica-major order.
END
}
in_arg {
name: "broadcast_inputs"
description: <<END
additional arguments to broadcast to all replicas. The
broadcast inputs are appended to the per-replica inputs when calling
computation.
END
}
in_arg {
name: "guaranteed_constants"
description: <<END
arguments which have been guaranteed to not
change their values during the session lifetime. These contain tensors marked as
constant using the GuaranteeConstOp.
END
}
out_arg {
name: "outputs"
description: <<END
the outputs of 'computation'.
END
}
attr {
name: "computation"
description: <<END
a function containing the computation to run.
END
}
attr {
name: "num_replicas"
description: <<END
the number of replicas of the computation to run.
END
}
attr {
name: "num_cores_per_replica"
description: <<END
the number of logical cores in each replica.
END
}
attr {
name: "topology"
description: <<END
A serialized tensorflow.tpu.TopologyProto that describes the TPU
topology.
END
}
attr {
name: "use_tpu"
description: <<END
a bool indicating if this computation will run on TPU or CPU/GPU.
Currently, only supports a default placement (computation is placed on GPU
if one is available, and on CPU if not).
END
}
attr {
name: "device_assignment"
description: <<END
a flattened array with shape
[replica, num_cores_per_replica, mesh_dimension] that maps the coordinates
of logical cores in each replica of a computation to physical coordinates in
the TPU topology.
END
}
attr {
name: "Tinputs"
description: <<END
the types of the arguments to 'computation'.
END
}
attr {
name: "Tbroadcast_inputs"
description: <<END
the types of the additional arguments to broadcast to all
replicas.
END
}
attr {
name: "Tguaranteed_constants"
description: <<END
the types of the arguments to 'guaranteed_constants'.
END
}
attr {
name: "output_types"
description: <<END
the types of the outputs of 'computation'.
END
}
summary: "Runs replicated computations on a distributed TPU system."
}

217
tensorflow/core/ops/compat/ops_history.v1.pbtxt
View File

@ -84403,223 +84403,6 @@ op {
type: "func"
}
}
op {
name: "TPUReplicate"
input_arg {
name: "inputs"
type_list_attr: "Tinputs"
}
input_arg {
name: "broadcast_inputs"
type_list_attr: "Tbroadcast_inputs"
}
input_arg {
name: "variables"
type: DT_RESOURCE
number_attr: "NumVariables"
}
input_arg {
name: "guaranteed_constants"
type_list_attr: "Tguaranteed_constants"
}
output_arg {
name: "outputs"
type_list_attr: "output_types"
}
attr {
name: "computation"
type: "func"
}
attr {
name: "num_replicas"
type: "int"
has_minimum: true
minimum: 1
}
attr {
name: "num_cores_per_replica"
type: "int"
default_value {
i: 1
}
}
attr {
name: "topology"
type: "string"
default_value {
s: ""
}
}
attr {
name: "use_tpu"
type: "bool"
default_value {
b: true
}
}
attr {
name: "device_assignment"
type: "list(int)"
default_value {
list {
}
}
}
attr {
name: "host_compute_core"
type: "list(string)"
default_value {
list {
}
}
}
attr {
name: "Tinputs"
type: "list(type)"
has_minimum: true
}
attr {
name: "Tbroadcast_inputs"
type: "list(type)"
has_minimum: true
}
attr {
name: "NumVariables"
type: "int"
has_minimum: true
}
attr {
name: "Tguaranteed_constants"
type: "list(type)"
has_minimum: true
}
attr {
name: "output_types"
type: "list(type)"
has_minimum: true
}
attr {
name: "padding_map"
type: "list(string)"
default_value {
list {
}
}
}
is_stateful: true
}
op {
name: "TPUReplicate"
input_arg {
name: "inputs"
type_list_attr: "Tinputs"
}
input_arg {
name: "broadcast_inputs"
type_list_attr: "Tbroadcast_inputs"
}
input_arg {
name: "variables"
type: DT_RESOURCE
number_attr: "NumVariables"
}
input_arg {
name: "guaranteed_constants"
type_list_attr: "Tguaranteed_constants"
}
output_arg {
name: "outputs"
type_list_attr: "output_types"
}
attr {
name: "computation"
type: "func"
}
attr {
name: "num_replicas"
type: "int"
has_minimum: true
minimum: 1
}
attr {
name: "num_cores_per_replica"
type: "int"
default_value {
i: 1
}
}
attr {
name: "topology"
type: "string"
default_value {
s: ""
}
}
attr {
name: "use_tpu"
type: "bool"
default_value {
b: true
}
}
attr {
name: "device_assignment"
type: "list(int)"
default_value {
list {
}
}
}
attr {
name: "host_compute_core"
type: "list(string)"
default_value {
list {
}
}
}
attr {
name: "Tinputs"
type: "list(type)"
has_minimum: true
}
attr {
name: "Tbroadcast_inputs"
type: "list(type)"
has_minimum: true
}
attr {
name: "NumVariables"
type: "int"
has_minimum: true
}
attr {
name: "Tguaranteed_constants"
type: "list(type)"
has_minimum: true
}
attr {
name: "output_types"
type: "list(type)"
has_minimum: true
}
attr {
name: "padding_map"
type: "list(string)"
default_value {
list {
}
}
}
attr {
name: "step_marker_location"
type: "string"
default_value {
s: "STEP_MARK_AT_ENTRY"
}
}
is_stateful: true
}
op {
name: "TPUReplicateMetadata"
attr {

217
tensorflow/core/ops/compat/ops_history.v2.pbtxt
View File

@ -82098,223 +82098,6 @@ op {
type: "func"
}
}
op {
name: "TPUReplicate"
input_arg {
name: "inputs"
type_list_attr: "Tinputs"
}
input_arg {
name: "broadcast_inputs"
type_list_attr: "Tbroadcast_inputs"
}
input_arg {
name: "variables"
type: DT_RESOURCE
number_attr: "NumVariables"
}
input_arg {
name: "guaranteed_constants"
type_list_attr: "Tguaranteed_constants"
}
output_arg {
name: "outputs"
type_list_attr: "output_types"
}
attr {
name: "computation"
type: "func"
}
attr {
name: "num_replicas"
type: "int"
has_minimum: true
minimum: 1
}
attr {
name: "num_cores_per_replica"
type: "int"
default_value {
i: 1
}
}
attr {
name: "topology"
type: "string"
default_value {
s: ""
}
}
attr {
name: "use_tpu"
type: "bool"
default_value {
b: true
}
}
attr {
name: "device_assignment"
type: "list(int)"
default_value {
list {
}
}
}
attr {
name: "host_compute_core"
type: "list(string)"
default_value {
list {
}
}
}
attr {
name: "Tinputs"
type: "list(type)"
has_minimum: true
}
attr {
name: "Tbroadcast_inputs"
type: "list(type)"
has_minimum: true
}
attr {
name: "NumVariables"
type: "int"
has_minimum: true
}
attr {
name: "Tguaranteed_constants"
type: "list(type)"
has_minimum: true
}
attr {
name: "output_types"
type: "list(type)"
has_minimum: true
}
attr {
name: "padding_map"
type: "list(string)"
default_value {
list {
}
}
}
is_stateful: true
}
op {
name: "TPUReplicate"
input_arg {
name: "inputs"
type_list_attr: "Tinputs"
}
input_arg {
name: "broadcast_inputs"
type_list_attr: "Tbroadcast_inputs"
}
input_arg {
name: "variables"
type: DT_RESOURCE
number_attr: "NumVariables"
}
input_arg {
name: "guaranteed_constants"
type_list_attr: "Tguaranteed_constants"
}
output_arg {
name: "outputs"
type_list_attr: "output_types"
}
attr {
name: "computation"
type: "func"
}
attr {
name: "num_replicas"
type: "int"
has_minimum: true
minimum: 1
}
attr {
name: "num_cores_per_replica"
type: "int"
default_value {
i: 1
}
}
attr {
name: "topology"
type: "string"
default_value {
s: ""
}
}
attr {
name: "use_tpu"
type: "bool"
default_value {
b: true
}
}
attr {
name: "device_assignment"
type: "list(int)"
default_value {
list {
}
}
}
attr {
name: "host_compute_core"
type: "list(string)"
default_value {
list {
}
}
}
attr {
name: "Tinputs"
type: "list(type)"
has_minimum: true
}
attr {
name: "Tbroadcast_inputs"
type: "list(type)"
has_minimum: true
}
attr {
name: "NumVariables"
type: "int"
has_minimum: true
}
attr {
name: "Tguaranteed_constants"
type: "list(type)"
has_minimum: true
}
attr {
name: "output_types"
type: "list(type)"
has_minimum: true
}
attr {
name: "padding_map"
type: "list(string)"
default_value {
list {
}
}
}
attr {
name: "step_marker_location"
type: "string"
default_value {
s: "STEP_MARK_AT_ENTRY"
}
}
is_stateful: true
}
op {
name: "TPUReplicateMetadata"
attr {

2
tensorflow/core/ops/tpu_replication_ops.cc
View File

@ -88,7 +88,7 @@ REGISTER_OP("TPUCompilationResult")
.Output("output: string")
.SetShapeFn(shape_inference::ScalarShape);
REGISTER_OP("TPUReplicate")
REGISTER_OP("_TPUReplicate")
.Attr("computation: func")
.Attr("num_replicas: int >= 1")
.Attr("num_cores_per_replica: int = 1")

4
tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
View File

@ -3980,10 +3980,6 @@ tf_module {
name: "TPUPartitionedCall"
argspec: "args=[\'args\', \'device_ordinal\', \'Tout\', \'f\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "TPUReplicate"
argspec: "args=[\'inputs\', \'broadcast_inputs\', \'variables\', \'guaranteed_constants\', \'computation\', \'num_replicas\', \'output_types\', \'num_cores_per_replica\', \'topology\', \'use_tpu\', \'device_assignment\', \'host_compute_core\', \'padding_map\', \'step_marker_location\', \'name\'], varargs=None, keywords=None, defaults=[\'1\', \'\', \'True\', \'[]\', \'[]\', \'[]\', \'STEP_MARK_AT_ENTRY\', \'None\'], "
}
member_method {
name: "TPUReplicateMetadata"
argspec: "args=[\'num_replicas\', \'num_cores_per_replica\', \'topology\', \'use_tpu\', \'device_assignment\', \'computation_shape\', \'host_compute_core\', \'padding_map\', \'step_marker_location\', \'name\'], varargs=None, keywords=None, defaults=[\'1\', \'\', \'True\', \'[]\', \'[]\', \'[]\', \'[]\', \'STEP_MARK_AT_ENTRY\', \'None\'], "

4
tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
View File

@ -3980,10 +3980,6 @@ tf_module {
name: "TPUPartitionedCall"
argspec: "args=[\'args\', \'device_ordinal\', \'Tout\', \'f\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "TPUReplicate"
argspec: "args=[\'inputs\', \'broadcast_inputs\', \'variables\', \'guaranteed_constants\', \'computation\', \'num_replicas\', \'output_types\', \'num_cores_per_replica\', \'topology\', \'use_tpu\', \'device_assignment\', \'host_compute_core\', \'padding_map\', \'step_marker_location\', \'name\'], varargs=None, keywords=None, defaults=[\'1\', \'\', \'True\', \'[]\', \'[]\', \'[]\', \'STEP_MARK_AT_ENTRY\', \'None\'], "
}
member_method {
name: "TPUReplicateMetadata"
argspec: "args=[\'num_replicas\', \'num_cores_per_replica\', \'topology\', \'use_tpu\', \'device_assignment\', \'computation_shape\', \'host_compute_core\', \'padding_map\', \'step_marker_location\', \'name\'], varargs=None, keywords=None, defaults=[\'1\', \'\', \'True\', \'[]\', \'[]\', \'[]\', \'[]\', \'STEP_MARK_AT_ENTRY\', \'None\'], "