Enable constant folding in L0 optimization level.

Change: 118861866

1 files changed, 15 insertions, 0 deletions

diff --git a/tensorflow/core/common_runtime/memory_types.cc b/tensorflow/core/common_runtime/memory_types.cc
index d8585859b8..957a124934 100644
--- a/tensorflow/core/common_runtime/memory_types.cc
+++ b/tensorflow/core/common_runtime/memory_types.cc
@@ -82,4 +82,19 @@ Status ValidateMemoryTypes(DeviceType device_type, const Graph* g) {
   return Status::OK();
 }
 
+Status MemoryTypeForOutput(DeviceType device_type, const Graph* g,
+                           const Node* n, int index, MemoryType* memory_type) {
+  MemoryTypeVector inp_mvec;
+  MemoryTypeVector out_mvec;
+  TF_RETURN_IF_ERROR(MemoryTypesForNode(g->op_registry(), device_type, n->def(),
+                                        &inp_mvec, &out_mvec));
+  if (out_mvec.size() <= index) {
+    return errors::Internal("Trying to get the memory type for ", index,
+                            "'th output of node ", n->DebugString(),
+                            " that has only ", out_mvec.size(), " outputs");
+  }
+  *memory_type = out_mvec[index];
+  return Status::OK();
+}
+
 }  // end namespace tensorflow