Add DataFormatVecPermute op.

PiperOrigin-RevId: 178421287

1 files changed, 17 insertions, 0 deletions

diff --git a/tensorflow/core/ops/nn_ops.cc b/tensorflow/core/ops/nn_ops.cc
index f58425db0a..8c31be0c0d 100644
--- a/tensorflow/core/ops/nn_ops.cc
+++ b/tensorflow/core/ops/nn_ops.cc
@@ -768,6 +768,23 @@ src_format: source data format.
 dst_format: destination data format.
 )doc");
 
+REGISTER_OP("DataFormatVecPermute")
+    .Input("x: T")
+    .Output("y: T")
+    .Attr("T: {int32, int64} = DT_INT32")
+    .Attr("src_format: string = 'NHWC'")
+    .Attr("dst_format: string = 'NCHW'")
+    .SetShapeFn(shape_inference::UnchangedShape)
+    .Doc(R"doc(
+Returns the permuted vector in the destination data format given the one in
+the source data format.
+
+x: Vector in source data format. Must be of size 4.
+y: Vector in destination data format. Must be of size 4.
+src_format: source data format.
+dst_format: destination data format.
+)doc");
+
 REGISTER_OP("FusedResizeAndPadConv2D")
     .Input("input: T")
     .Input("size: int32")