Adds a tf.ensure_shape function as a substitute for tensor.set_shape, which validates the true shape of the tensor at runtime.

PiperOrigin-RevId: 210570878

2 files changed, 30 insertions, 0 deletions

diff --git a/tensorflow/core/api_def/base_api/api_def_EnsureShape.pbtxt b/tensorflow/core/api_def/base_api/api_def_EnsureShape.pbtxt
new file mode 100644
index 0000000000..1658472209
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_EnsureShape.pbtxt
@@ -0,0 +1,26 @@
+op {
+  graph_op_name: "EnsureShape"
+  in_arg {
+  name: "input"
+  description: <<END
+A tensor, whose shape is to be validated.
+END
+  }
+  out_arg {
+    name: "output"
+    description: <<END
+A tensor with the same shape and contents as the input tensor or value.
+END
+  }
+  attr {
+    name: "shape"
+    description: <<END
+The expected (possibly partially specified) shape of the input tensor.
+END
+  }
+  summary: "Ensures that the tensor's shape matches the expected shape."
+  description: <<END
+Raises an error if the input tensor's shape does not match the specified shape.
+Returns the input tensor otherwise.
+END
+}
diff --git a/tensorflow/core/api_def/python_api/api_def_EnsureShape.pbtxt b/tensorflow/core/api_def/python_api/api_def_EnsureShape.pbtxt
new file mode 100644
index 0000000000..4414d973ac
--- /dev/null
+++ b/tensorflow/core/api_def/python_api/api_def_EnsureShape.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "EnsureShape"
+  visibility: HIDDEN
+}