chore: import upstream snapshot with attribution

tensorflow/cc/ops/const_op.cc

@@ -0,0 +1,86 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/core/framework/types.h"
+
+namespace tensorflow {
+namespace ops {
+
+namespace {
+template <typename T>
+Output ConstHelper(const Scope& scope, const T& value, DataType dtype) {
+  if (!scope.ok()) return Output();
+
+  Node* ret;
+  Graph* graph = scope.graph();
+  const std::string unique_name = scope.GetUniqueNameForOp("Const");
+  auto builder = NodeBuilder(unique_name, "Const")
+                     .Attr("value", value)
+                     .Attr("dtype", dtype);
+  scope.UpdateBuilder(&builder);
+  scope.UpdateStatus(builder.Finalize(graph, &ret));
+  if (!scope.ok()) return Output();
+
+  scope.UpdateStatus(scope.DoShapeInference(ret));
+  if (!scope.ok()) return Output();
+
+  return Output(ret);
+}
+}  // namespace
+
+Output Const(const Scope& scope, const Input::Initializer& val) {
+  if (!val.status.ok()) {
+    scope.UpdateStatus(val.status);
+    return Output();
+  }
+  return ConstHelper(scope, val.tensor, val.tensor.dtype());
+}
+
+Output ConstFromProto(const Scope& scope, const TensorProto& proto) {
+  return ConstHelper(scope, proto, proto.dtype());
+}
+
+NodeBuilder::NodeOut AsNodeOut(const Scope& scope, const Input& inp) {
+  if (!inp.status().ok()) {
+    scope.UpdateStatus(inp.status());
+    return NodeBuilder::NodeOut(inp.node(), inp.index());
+  }
+  if (inp.node()) {
+    return NodeBuilder::NodeOut(inp.node(), inp.index());
+  }
+  if (!inp.node_name().empty()) {
+    return NodeBuilder::NodeOut(inp.node_name(), inp.index(), inp.data_type());
+  }
+  auto transformed = Input{
+      Const(scope.NewSubScope("Const"), Input::Initializer(inp.tensor()))};
+  return NodeBuilder::NodeOut{transformed.node(), transformed.index()};
+}
+
+std::vector<NodeBuilder::NodeOut> AsNodeOutList(const Scope& scope,
+                                                const InputList& inp) {
+  std::vector<NodeBuilder::NodeOut> out;
+  for (const auto& i : inp) {
+    const auto node_out = AsNodeOut(scope, i);
+    if (!scope.ok()) {
+      return {};
+    }
+    out.push_back(node_out);
+  }
+  return out;
+}
+
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/ops/const_op.h

@@ -0,0 +1,87 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_OPS_CONST_OP_H_
+#define TENSORFLOW_CC_OPS_CONST_OP_H_
+
+#include <vector>
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/core/graph/node_builder.h"
+
+namespace tensorflow {
+namespace ops {
+
+/// @defgroup const_op Const Op
+/// @{
+
+Output Const(const Scope& scope, const Input::Initializer& val);
+
+Output ConstFromProto(const Scope& scope, const TensorProto& proto);
+
+NodeBuilder::NodeOut AsNodeOut(const Scope& scope, const Input& inp);
+
+template <typename T>
+Output Const(const Scope& scope, const Input::Initializer& val) {
+  auto orig_const_output = Const(scope, val);
+  if (!scope.ok()) return Output();
+
+  typedef typename Input::Initializer::RealType<T>::type DstT;
+
+  if (val.tensor.dtype() == DataTypeToEnum<DstT>::v()) {
+    return orig_const_output;
+  }
+  if (val.tensor.NumElements() == 0) {
+    Tensor t(DataTypeToEnum<DstT>::v(), val.tensor.shape());
+    return Const(scope, Input::Initializer(t));
+  }
+
+  // TODO(keveman): Refactor Cast op's kernel implementation such that the code
+  // can be directly called here instead of adding the Cast op to the graph.
+  auto orig_const = AsNodeOut(scope, orig_const_output);
+  const auto cast_op_name = scope.GetUniqueNameForOp("Cast");
+
+  auto cast_builder = NodeBuilder(cast_op_name, "Cast")
+                          .Input(orig_const)
+                          .Attr("DstT", DataTypeToEnum<DstT>::v());
+  scope.UpdateBuilder(&cast_builder);
+  Node* ret;
+  scope.UpdateStatus(cast_builder.Finalize(scope.graph(), &ret));
+  if (!scope.ok()) return Output();
+  scope.UpdateStatus(scope.DoShapeInference(ret));
+  return Output(ret, 0);
+}
+
+template <typename T>
+Output Const(const Scope& scope, const T& v, const TensorShape shape) {
+  return Const(scope, Input::Initializer(v, shape));
+}
+
+template <typename T>
+Output Const(const Scope& scope, const std::initializer_list<T>& v,
+             const TensorShape shape) {
+  return Const(scope, Input::Initializer(v, shape));
+}
+
+std::vector<NodeBuilder::NodeOut> AsNodeOutList(const Scope& scope,
+                                                const InputList& inp);
+
+/// }@
+
+}  // namespace ops
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_OPS_CONST_OP_H_

tensorflow/cc/ops/const_op_test.cc

@@ -0,0 +1,151 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+
+namespace {
+
+template <typename T>
+void ExpectNodeEqual(const Node* n, gtl::ArraySlice<T> values,
+                     TensorShape shape) {
+  EXPECT_TRUE(n->IsConstant());
+  Tensor tensor;
+  TF_EXPECT_OK(GetNodeAttr(n->attrs(), "value", &tensor));
+  DataType dtype;
+  TF_EXPECT_OK(GetNodeAttr(n->attrs(), "dtype", &dtype));
+  EXPECT_EQ(tensor.dtype(), dtype);
+  test::ExpectTensorEqual<T>(tensor, test::AsTensor(values, shape));
+}
+
+void ExpectTypeAndShape(const Node* n, DataType expected_dtype,
+                        TensorShape expected_shape) {
+  EXPECT_TRUE(n->IsConstant());
+  Tensor tensor;
+  TF_EXPECT_OK(GetNodeAttr(n->attrs(), "value", &tensor));
+  DataType dtype;
+  TF_EXPECT_OK(GetNodeAttr(n->attrs(), "dtype", &dtype));
+  EXPECT_EQ(dtype, expected_dtype);
+  EXPECT_EQ(expected_shape, TensorShape(tensor.shape()));
+}
+
+}  // namespace
+
+TEST(ConstOpTest, Basic) {
+  Scope root = Scope::NewRootScope();
+  auto c = ops::Const(root, 42.0f);
+  TF_EXPECT_OK(root.status());
+  EXPECT_EQ(c.op().output_type(0), DT_FLOAT);
+  ExpectNodeEqual<float>(c.node(), {42.0f}, {});
+}
+
+TEST(ConstOpTest, MultiDim) {
+  Scope root = Scope::NewRootScope();
+  auto c = ops::Const(root, {{2.0}, {3.0}});
+  TF_CHECK_OK(root.status());
+  EXPECT_EQ(c.op().output_type(0), DT_DOUBLE);
+  ExpectNodeEqual<double>(c.node(), {2.0, 3.0}, {2, 1});
+}
+
+TEST(ConstOpTest, Empty) {
+  Scope root = Scope::NewRootScope();
+
+  auto c1 = ops::Const(root, {});
+  TF_CHECK_OK(root.status());
+  ExpectTypeAndShape(c1.node(), DT_FLOAT, {0});
+
+  auto c2 = ops::Const(root, {{}});
+  TF_CHECK_OK(root.status());
+  ExpectTypeAndShape(c2.node(), DT_FLOAT, {1, 0});
+
+  auto c3 = ops::Const(root, {{{}, {}}});
+  TF_CHECK_OK(root.status());
+  ExpectTypeAndShape(c3.node(), DT_FLOAT, {1, 2, 0});
+
+  auto c4 = ops::Const<int>(root, {{{}}});
+  TF_CHECK_OK(root.status());
+  ExpectTypeAndShape(c4.node(), DT_INT32, {1, 1, 0});
+
+  ops::Const(root, {{}, {{}}});
+  EXPECT_FALSE(root.status().ok());
+}
+
+TEST(ConstOpTest, WithExplicitShape) {
+  Scope root = Scope::NewRootScope();
+  auto c = ops::Const(root, 42.0, {2, 2});
+  TF_CHECK_OK(root.status());
+  EXPECT_EQ(c.op().output_type(0), DT_DOUBLE);
+  ExpectNodeEqual<double>(c.node(), {42.0, 42.0, 42.0, 42.0}, {2, 2});
+
+  auto d = ops::Const(root, {"1", "2", "3", "4", "5", "6"}, {2, 3});
+  TF_CHECK_OK(root.status());
+  EXPECT_EQ(d.op().output_type(0), DT_STRING);
+  ExpectNodeEqual<tstring>(d.node(), {"1", "2", "3", "4", "5", "6"}, {2, 3});
+}
+
+TEST(ConstOpTest, FromProto) {
+  Scope root = Scope::NewRootScope();
+  TensorProto proto;
+  proto.set_dtype(DT_DOUBLE);
+  TensorShape({2, 2}).AsProto(proto.mutable_tensor_shape());
+  for (int i = 0; i < 4; ++i) {
+    proto.add_double_val(static_cast<double>(i));
+  }
+  auto c = ops::ConstFromProto(root, proto);
+  TF_CHECK_OK(root.status());
+  EXPECT_EQ(c.op().output_type(0), DT_DOUBLE);
+  ExpectNodeEqual<double>(c.node(), {0.0, 1.0, 2.0, 3.0}, {2, 2});
+}
+
+TEST(ConstOpTest, InvalidInitializer) {
+  Scope root = Scope::NewRootScope();
+  ops::Const(root, {{2.0}, {"df"}});
+  EXPECT_FALSE(root.status().ok());
+}
+
+TEST(ConstOpTest, Names) {
+  Scope root = Scope::NewRootScope();
+  auto c = ops::Const(root, {{2.0}, {3.0}});
+  EXPECT_EQ(c.node()->name(), "Const");
+  auto c_1 = ops::Const(root, {{2.0}, {3.0}});
+  EXPECT_EQ(c_1.node()->name(), "Const_1");
+
+  auto x = ops::Const(root.WithOpName("x"), 1);
+  EXPECT_EQ(x.node()->name(), "x");
+  auto x_1 = ops::Const(root.WithOpName("x"), 1);
+  EXPECT_EQ(x_1.node()->name(), "x_1");
+
+  Scope child = root.NewSubScope("c");
+  auto c_y = ops::Const(child.WithOpName("y"), 1);
+  EXPECT_EQ(c_y.node()->name(), "c/y");
+  auto c_y_1 = ops::Const(child.WithOpName("y"), 1);
+  EXPECT_EQ(c_y_1.node()->name(), "c/y_1");
+}
+
+TEST(ConstOpTest, TemplatedConst) {
+  Scope root = Scope::NewRootScope();
+  auto c1 = ops::Const<int>(root, {1, 2});
+  ExpectTypeAndShape(c1.node(), DT_INT32, {2});
+
+  auto c2 = ops::Const<tstring>(root, {{"this"}, {"is"}, {"a"}, {"constant"}});
+  ExpectTypeAndShape(c2.node(), DT_STRING, {4, 1});
+}
+
+}  // namespace tensorflow

tensorflow/cc/ops/standard_ops.h

@@ -0,0 +1,40 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_OPS_STANDARD_OPS_H_
+#define TENSORFLOW_CC_OPS_STANDARD_OPS_H_
+
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/candidate_sampling_ops.h"
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/cc/ops/control_flow_ops.h"
+#include "tensorflow/cc/ops/data_flow_ops.h"
+#include "tensorflow/cc/ops/image_ops.h"
+#include "tensorflow/cc/ops/io_ops.h"
+#include "tensorflow/cc/ops/linalg_ops.h"
+#include "tensorflow/cc/ops/logging_ops.h"
+#include "tensorflow/cc/ops/lookup_ops.h"
+#include "tensorflow/cc/ops/math_ops.h"
+#include "tensorflow/cc/ops/nn_ops.h"
+#include "tensorflow/cc/ops/no_op.h"
+#include "tensorflow/cc/ops/parsing_ops.h"
+#include "tensorflow/cc/ops/random_ops.h"
+#include "tensorflow/cc/ops/sparse_ops.h"
+#include "tensorflow/cc/ops/state_ops.h"
+#include "tensorflow/cc/ops/string_ops.h"
+#include "tensorflow/cc/ops/training_ops.h"
+#include "tensorflow/cc/ops/user_ops.h"
+
+#endif  // TENSORFLOW_CC_OPS_STANDARD_OPS_H_

tensorflow/cc/ops/while_loop.cc

@@ -0,0 +1,251 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/ops/while_loop.h"
+
+#include "tensorflow/cc/framework/scope_internal.h"
+#include "tensorflow/cc/ops/control_flow_ops_internal.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/common_runtime/shape_refiner.h"
+#include "tensorflow/core/graph/node_builder.h"
+
+namespace tensorflow {
+namespace ops {
+
+namespace {
+
+// Utility function for converting to internal C++ datatypes.
+OutputTensor ToOutputTensor(const Output& output) {
+  return OutputTensor(output.node(), output.index());
+}
+
+// Utility function for converting to internal C++ datatypes.
+std::vector<OutputTensor> ToOutputTensors(const std::vector<Output>& outputs) {
+  std::vector<OutputTensor> result(outputs.size());
+  for (int i = 0; i < outputs.size(); ++i) {
+    result[i] = ToOutputTensor(outputs[i]);
+  }
+  return result;
+}
+
+// Utility function for converting to internal C++ datatypes.
+std::vector<Node*> ToNodes(const std::vector<Output>& outputs) {
+  std::vector<Node*> result(outputs.size());
+  for (int i = 0; i < outputs.size(); ++i) {
+    result[i] = outputs[i].node();
+  }
+  return result;
+}
+
+// Manually generates the name of the `loop_var_idx`-th NextIteration node of a
+// loop being constructed with `scope`. This is used to define the backedge
+// before the NextIteration node is created.
+std::string NextIterationName(const Scope& scope, int loop_var_idx) {
+  std::string result;
+  const std::string& prefix = scope.impl()->name();
+  if (!prefix.empty()) absl::StrAppend(&result, prefix, "/");
+  absl::StrAppend(&result, "NextIteration");
+  if (loop_var_idx > 0) absl::StrAppend(&result, "_", loop_var_idx);
+  return result;
+}
+
+// Creates the `loop_var_idx`-th Merge node of a loop being constructed with
+// `scope`. `enter_output` is the `loop_var_idx`-th Enter node's output.
+absl::Status CreateMerge(const Scope& scope, int loop_var_idx,
+                         const Output& enter_output, Output* merge_output) {
+  // The merge nodes accept the while loop's back edges as an input (i.e. the
+  // not-yet-created next iteration nodes). Use the underlying NodeBuilder API
+  // directly to create the back edge.
+  NodeBuilder::NodeOut enter_input(enter_output.node(), enter_output.index());
+
+  const int next_output_index = 0;
+  DataType dtype = enter_output.node()->output_type(0);
+  NodeBuilder::NodeOut next_input(NextIterationName(scope, loop_var_idx),
+                                  next_output_index, dtype);
+
+  std::vector<NodeBuilder::NodeOut> input_list({enter_input, next_input});
+  const std::string unique_name = scope.GetUniqueNameForOp("Merge");
+  NodeBuilder builder = NodeBuilder(unique_name, "Merge").Input(input_list);
+  scope.UpdateBuilder(&builder);
+
+  Node* merge_node;
+  TF_RETURN_IF_ERROR(builder.Finalize(scope.graph(), &merge_node));
+  TF_RETURN_IF_ERROR(scope.DoShapeInference(merge_node));
+  *merge_output = Output(merge_node, 0);
+  return absl::OkStatus();
+}
+
+// Creates the condition subgraph defined by `cond`.
+absl::Status CreateCond(const Scope& scope, const CondGraphBuilderFn& cond,
+                        const std::vector<Output>& inputs, Output* output) {
+  // The control dependency is for constants in the cond graph, and other ops
+  // that do not depend on the loop variables. This ensures that these ops are
+  // in the while loop frame (since they will indirectly depend on an Enter node
+  // defining the frame) and that they are executed once per loop iteration.
+  //
+  // TODO(skyewm): the control dep will be added to all nodes in the cond graph.
+  // This is at best unnecessary, and at worst may prevent different parts of
+  // different loop iterations from executing in parallel.
+  Scope cond_scope =
+      scope.NewSubScope("cond").WithControlDependencies(inputs[0]);
+  Output raw_cond_out;
+  TF_RETURN_IF_ERROR(cond(cond_scope, inputs, &raw_cond_out));
+
+  TF_RETURN_IF_ERROR(scope.graph()->IsValidOutputTensor(raw_cond_out.node(),
+                                                        raw_cond_out.index()));
+  if (raw_cond_out.type() != DT_BOOL) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "BuildWhileLoop: 'cond' argument must return a boolean output, got ",
+        DataTypeString(raw_cond_out.type())));
+  }
+  // TODO(skyewm): check that raw_cond_out is scalar
+
+  *output = LoopCond(scope, raw_cond_out).output;
+  return absl::OkStatus();
+}
+
+// Create the body subgraph defined by `body`. `outputs` must be non-null and
+// empty.
+absl::Status CreateBody(const Scope& scope, const BodyGraphBuilderFn& body,
+                        const std::vector<Output>& inputs,
+                        std::vector<Output>* outputs) {
+  DCHECK(outputs != nullptr);
+  DCHECK(outputs->empty());
+
+  // The control dependency is analogous to that in CreateCond().
+  Scope body_scope =
+      scope.NewSubScope("body").WithControlDependencies(inputs[0]);
+  TF_RETURN_IF_ERROR(body(body_scope, inputs, outputs));
+
+  const size_t num_loop_vars = inputs.size();
+  if (outputs->size() != num_loop_vars) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("BuildWhileLoop: 'body' argument expected to return ",
+                     num_loop_vars, " output(s), got ", outputs->size()));
+  }
+  for (const Output& output : *outputs) {
+    TF_RETURN_IF_ERROR(
+        scope.graph()->IsValidOutputTensor(output.node(), output.index()));
+    // TODO(skyewm): check output types/shapes
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+// A while loop with a single loop variable looks like this:
+//
+// (output)
+//     ^    +---------------+
+//     |    | body subgraph +-------------+
+//    Exit  +---------------+             |
+//      ^    ^                            |
+//      |    |                            |
+//      Switch<--------+                  v
+//        ^            |             NextIteration
+//        |     +------+--------+         |
+//        +---->| cond subgraph |         |
+//        |     +---------------+         |
+//       Merge<---------------------------+
+//       ^
+//       |
+//    Enter
+//      ^
+//      |
+//   (input)
+//
+// If there are multiple loop variables, each of the control flow ops is
+// duplicated for each loop variable.
+// TODO(skyewm): link to public version of design doc
+absl::Status BuildWhileLoop(const Scope& scope,
+                            const std::vector<Output>& inputs,
+                            const CondGraphBuilderFn& cond,
+                            const BodyGraphBuilderFn& body,
+                            const std::string& frame_name, OutputList* outputs,
+                            bool create_while_ctx, Output* cond_output) {
+  DCHECK(!inputs.empty());
+  DCHECK(outputs != nullptr);
+  DCHECK(outputs->empty());
+
+  TF_RETURN_IF_ERROR(scope.status());
+  const size_t num_loop_vars = inputs.size();
+
+  std::vector<Output> enter_outputs(num_loop_vars);
+  for (size_t i = 0; i < num_loop_vars; ++i) {
+    enter_outputs[i] = internal::Enter(scope, inputs[i], frame_name);
+  }
+  TF_RETURN_IF_ERROR(scope.status());
+
+  std::vector<Output> merge_outputs(num_loop_vars);
+  for (size_t i = 0; i < num_loop_vars; ++i) {
+    TF_RETURN_IF_ERROR(
+        CreateMerge(scope, i, enter_outputs[i], &merge_outputs[i]));
+  }
+
+  Output cond_out;
+  TF_RETURN_IF_ERROR(CreateCond(scope, cond, merge_outputs, &cond_out));
+  if (cond_output != nullptr) *cond_output = cond_out;
+
+  std::vector<Output> switch_trues(num_loop_vars);
+  std::vector<Output> switch_falses(num_loop_vars);
+  for (size_t i = 0; i < num_loop_vars; ++i) {
+    auto switch_i = Switch(scope, merge_outputs[i], cond_out);
+    switch_trues[i] = switch_i.output_true;
+    switch_falses[i] = switch_i.output_false;
+  }
+  TF_RETURN_IF_ERROR(scope.status());
+
+  std::vector<Output> body_outputs;
+  TF_RETURN_IF_ERROR(CreateBody(scope, body, switch_trues, &body_outputs));
+
+  std::vector<Output> next_outputs(num_loop_vars);
+  for (size_t i = 0; i < num_loop_vars; ++i) {
+    next_outputs[i] = NextIteration(scope, body_outputs[i]);
+    DCHECK_EQ(next_outputs[i].node()->name(), NextIterationName(scope, i));
+  }
+  TF_RETURN_IF_ERROR(scope.status());
+
+  // Create the backedges from the NextIteration nodes to the Merge nodes.
+  for (size_t i = 0; i < num_loop_vars; ++i) {
+    const int merge_backedge_output_index = 1;
+    scope.graph()->AddEdge(next_outputs[i].node(), next_outputs[i].index(),
+                           merge_outputs[i].node(),
+                           merge_backedge_output_index);
+  }
+
+  outputs->resize(num_loop_vars);
+  for (size_t i = 0; i < num_loop_vars; ++i) {
+    (*outputs)[i] = internal::Exit(scope, switch_falses[i]);
+  }
+  TF_RETURN_IF_ERROR(scope.status());
+
+  if (create_while_ctx) {
+    WhileContext* while_ctx;
+    TF_RETURN_IF_ERROR(scope.graph()->AddWhileContext(
+        frame_name, ToNodes(enter_outputs), ToNodes(*outputs),
+        ToOutputTensor(cond_out), ToOutputTensors(switch_trues),
+        ToOutputTensors(body_outputs), &while_ctx));
+
+    // Set while_ctx for all exit nodes. We currently don't require knowing the
+    // while_ctx for any other nodes.
+    for (size_t i = 0; i < num_loop_vars; ++i) {
+      (*outputs)[i].node()->set_while_ctx(while_ctx);
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace ops
+}  // namespace tensorflow

tensorflow/cc/ops/while_loop.h

@@ -0,0 +1,80 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_OPS_WHILE_LOOP_H_
+#define TENSORFLOW_CC_OPS_WHILE_LOOP_H_
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+
+namespace tensorflow {
+namespace ops {
+
+// Function that takes cond graph inputs and returns cond graph boolean output.
+// 'output' need not be set if an error is returned.
+typedef std::function<absl::Status(
+    const Scope&, const std::vector<Output>& inputs, Output* output)>
+    CondGraphBuilderFn;
+
+// Function that takes body graph inputs and returns body graph outputs.
+// 'outputs' need not be populated if an error is returned.
+typedef std::function<absl::Status(const Scope&,
+                                   const std::vector<Output>& inputs,
+                                   std::vector<Output>* outputs)>
+    BodyGraphBuilderFn;
+
+// Constructs a while loop.
+//
+// Arguments:
+// * scope: used to construct the while loop.
+// * inputs: the initial values of the loop variables. Must be non-empty.
+// * cond: a function that builds the condition graph of the loop. Takes the
+//     current loop variables as inputs and returns a scalar boolean Output
+//     indicating whether the loop should continue.
+// * body: a function that builds the body graph of the loop. Takes the current
+//     loop variables as inputs and returns the updated loop variables.
+// * frame_name: the frame name to use for this while loop. This should be a
+//     unique name. This will be used as a prefix for created operations.
+// * outputs: output param that returns final loop variable outputs in non-error
+//     case. Must be non-null and empty.
+// * create_while_ctx: if true, a WhileContext is created and populated for this
+//     loop. See core/graph/while_context.h for more details on
+//     WhileContexts. This is set to false for loops used as part of gradient
+//     computations, since they're part of the gradient for a loop in the
+//     forward-pass.
+//     TODO(skyewm): revisit this. Should we create WhileContexts for all loops,
+//     even if we don't need them?
+// * cond_output: if non-null, the output of the predicate is returned. This
+//     will always be a LoopCond node.
+//
+// Returns an error if the while loop could not be fully constructed.
+//
+// TODO(skyewm): clean up partially-constructed loop in error case
+// TODO(skyewm): create public interface to this method
+absl::Status BuildWhileLoop(const Scope& scope,
+                            const std::vector<Output>& inputs,
+                            const CondGraphBuilderFn& cond,
+                            const BodyGraphBuilderFn& body,
+                            const std::string& frame_name, OutputList* outputs,
+                            bool create_while_ctx = true,
+                            Output* cond_output = nullptr);
+
+}  // namespace ops
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_OPS_WHILE_LOOP_H_

tensorflow/cc/ops/while_loop_test.cc

@@ -0,0 +1,201 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/ops/while_loop.h"
+#include "tensorflow/cc/client/client_session.h"
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/graph/while_context.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+
+namespace {
+
+class WhileLoopTest : public ::testing::Test {
+ protected:
+  WhileLoopTest() : scope_(Scope::NewRootScope()) {}
+
+  void Init(int num_inputs, DataType dtype = DT_INT32) {
+    for (int i = 0; i < num_inputs; ++i) {
+      inputs_.push_back(ops::Placeholder(scope_, dtype));
+    }
+  }
+
+  void CreateLoop(const ops::CondGraphBuilderFn& cond,
+                  const ops::BodyGraphBuilderFn& body,
+                  error::Code error_code = error::OK,
+                  const std::string& error_msg = "") {
+    absl::Status s =
+        ops::BuildWhileLoop(scope_, inputs_, cond, body, kFrameName, &outputs_);
+    EXPECT_EQ(s.code(), error_code);
+    EXPECT_EQ(s.message(), error_msg);
+  }
+
+  template <typename T>
+  void Run(const std::vector<Input::Initializer>& input_values,
+           const std::vector<T>& expected_output_values) {
+    ClientSession session(scope_);
+
+    DCHECK_EQ(input_values.size(), inputs_.size());
+    ClientSession::FeedType feeds;
+    for (int i = 0; i < inputs_.size(); ++i) {
+      feeds.emplace(inputs_[i], input_values[i]);
+    }
+
+    std::vector<Tensor> out_tensors;
+    TF_ASSERT_OK(session.Run(feeds, outputs_, &out_tensors));
+    ASSERT_EQ(out_tensors.size(), outputs_.size());
+
+    DCHECK_EQ(expected_output_values.size(), out_tensors.size());
+    for (int i = 0; i < out_tensors.size(); ++i) {
+      test::ExpectTensorEqual<T>(
+          out_tensors[i], test::AsTensor<T>({expected_output_values[i]}, {}));
+    }
+  }
+
+  Scope scope_;
+  std::vector<Output> inputs_;
+  std::vector<Output> outputs_;
+
+  static const char* const kFrameName;
+};
+
+const char* const WhileLoopTest::kFrameName = "test_loop";
+
+absl::Status LessThanTenCond(const Scope& s, const std::vector<Output>& inputs,
+                             Output* output) {
+  *output = ops::Less(s, inputs[0], 10);
+  return s.status();
+}
+
+absl::Status AddOneBody(const Scope& s, const std::vector<Output>& inputs,
+                        std::vector<Output>* outputs) {
+  outputs->push_back(ops::Add(s, inputs[0], 1));
+  return s.status();
+}
+
+TEST_F(WhileLoopTest, Basic) {
+  // Create loop: while (i < 10) i += 1
+  Init(1);
+  CreateLoop(LessThanTenCond, AddOneBody);
+
+  // Verify some output invariants
+  WhileContext* while_ctx;
+  for (int i = 0; i < outputs_.size(); ++i) {
+    Node* node = outputs_[i].node();
+    ASSERT_TRUE(node->IsExit()) << "Output node " << i << ":\n"
+                                << node->DebugString();
+    ASSERT_TRUE(node->while_ctx() != nullptr) << i;
+    if (i == 0) {
+      while_ctx = node->while_ctx();
+      EXPECT_EQ(while_ctx->frame_name(), kFrameName);
+    } else {
+      EXPECT_EQ(node->while_ctx(), while_ctx) << i;
+    }
+  }
+
+  // Run the loop and test we get the expected results
+  Run<int>({1}, {10});
+  Run<int>({11}, {11});
+}
+
+TEST_F(WhileLoopTest, WrongCondOutputType) {
+  Init(1);
+  CreateLoop(
+      [](const Scope& s, const std::vector<Output>& inputs, Output* output) {
+        *output = ops::Placeholder(s, DT_FLOAT);
+        return s.status();
+      },
+      AddOneBody, error::INVALID_ARGUMENT,
+      "BuildWhileLoop: 'cond' argument must return a boolean output, got "
+      "float");
+}
+
+// TODO(skyewm): test bad cond output shape
+
+TEST_F(WhileLoopTest, NullCondOutputNode) {
+  Init(1);
+  // TODO(skyewm): improve error message
+  CreateLoop(
+      [](const Scope& s, const std::vector<Output>& inputs, Output* output) {
+        *output = {nullptr, 0};
+        return s.status();
+      },
+      AddOneBody, error::INVALID_ARGUMENT, "Node is null");
+}
+
+TEST_F(WhileLoopTest, InvalidCondOutputIndex) {
+  Init(1);
+  CreateLoop(
+      [](const Scope& s, const std::vector<Output>& inputs, Output* output) {
+        auto less = ops::Less(s, inputs[0], 10);
+        *output = {less.node(), 100};
+        return s.status();
+      },
+      AddOneBody, error::OUT_OF_RANGE,
+      "Node 'cond/Less' (type: 'Less', num of outputs: 1) does not have output "
+      "100");
+}
+
+TEST_F(WhileLoopTest, UnsetCondOutput) {
+  Init(1);
+  CreateLoop([](const Scope& s, const std::vector<Output>& inputs,
+                Output* output) { return s.status(); },
+             AddOneBody, error::INVALID_ARGUMENT, "Node is null");
+}
+
+// TODO(skyewm): test bad body output type
+// TODO(skyewm): test bad body output shape
+
+TEST_F(WhileLoopTest, NullBodyOutputNode) {
+  Init(1);
+  // TODO(skyewm): improve error message
+  CreateLoop(LessThanTenCond,
+             [](const Scope& s, const std::vector<Output>& inputs,
+                std::vector<Output>* outputs) {
+               outputs->push_back({nullptr, 0});
+               return s.status();
+             },
+             error::INVALID_ARGUMENT, "Node is null");
+}
+
+TEST_F(WhileLoopTest, InvalidBodyOutputIndex) {
+  Init(1);
+  CreateLoop(LessThanTenCond,
+             [](const Scope& s, const std::vector<Output>& inputs,
+                std::vector<Output>* outputs) {
+               auto add = ops::Add(s, inputs[0], 1);
+               outputs->emplace_back(add.node(), 100);
+               return s.status();
+             },
+             error::OUT_OF_RANGE,
+             "Node 'body/Add' (type: 'Add', num of outputs: 1) does not have "
+             "output 100");
+}
+
+TEST_F(WhileLoopTest, UnsetBodyOutputs) {
+  Init(1);
+  CreateLoop(
+      LessThanTenCond,
+      [](const Scope& s, const std::vector<Output>& inputs,
+         std::vector<Output>* outputs) { return s.status(); },
+      error::INVALID_ARGUMENT,
+      "BuildWhileLoop: 'body' argument expected to return 1 output(s), got 0");
+}
+
+}  // namespace
+}  // namespace tensorflow