diff options
| author |  Rohan Jain <rohanj@google.com> | 2018-09-23 08:53:45 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-09-23 08:58:05 -0700 |
| commit | f1237459efb3a5578885b03d5b33c3fed350c348 (patch) | |
| tree | 128b49a2112e773640a06b3b0578afa8cddd81f0 /tensorflow/core/ops | |
| parent | 646b3c237deaddddd087d39ab57130b08375c4c7 (diff) | |
Moving MultiDeviceIterator from contrib to core.
PiperOrigin-RevId: 214173896
Diffstat (limited to 'tensorflow/core/ops')
| -rw-r--r-- | tensorflow/core/ops/dataset_ops.cc | 37 |
1 files changed, 37 insertions, 0 deletions
diff --git a/tensorflow/core/ops/dataset_ops.cc b/tensorflow/core/ops/dataset_ops.cc index 4d3f272c1b..1ada623cf5 100644 --- a/tensorflow/core/ops/dataset_ops.cc +++ b/tensorflow/core/ops/dataset_ops.cc @@ -932,4 +932,41 @@ REGISTER_OP("MapDefun") return Status::OK(); }); +REGISTER_OP("MultiDeviceIterator") + .Output("handle: resource") + .Attr("devices: list(string) >= 1") + .Attr("shared_name: string") + .Attr("container: string") + .Attr("output_types: list(type) >= 1") + .Attr("output_shapes: list(shape) >= 1") + .SetShapeFn(shape_inference::ScalarShape); + +REGISTER_OP("MultiDeviceIteratorInit") + .Input("dataset: variant") + .Input("multi_device_iterator: resource") + .Input("max_buffer_size: int64") + .Output("incarnation_id: int64") + .SetShapeFn(shape_inference::ScalarShape); + +REGISTER_OP("MultiDeviceIteratorGetNextFromShard") + .Input("multi_device_iterator: resource") + .Input("shard_num: int32") + .Input("incarnation_id: int64") + .Output("components: output_types") + .Attr("output_types: list(type) >= 1") + .Attr("output_shapes: list(shape) >= 1") + .SetShapeFn(IteratorGetNextShapeFn); + +REGISTER_OP("MultiDeviceIteratorToStringHandle") + .Input("multi_device_iterator: resource") + .Output("string_handle: string") + .SetShapeFn(shape_inference::ScalarShape); + +REGISTER_OP("MultiDeviceIteratorFromStringHandle") + .Input("string_handle: string") + .Output("multi_device_iterator: resource") + .Attr("output_types: list(type) >= 0 = []") + .Attr("output_shapes: list(shape) >= 0 = []") + .SetShapeFn(shape_inference::ScalarShape); + } // namespace tensorflow |