}
  }
}
node {
  name: "while/Less"
  op: "Less"
  input: "while/Merge"
  input: "while/Less/y"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
}
node {
  name: "while/LoopCond"
  op: "LoopCond"
  input: "while/Less"
}
node {
  name: "while/Switch"
  op: "Switch"
  input: "while/Merge"
  input: "while/LoopCond"
  attr {
    key: "T"

    %8:2 = tf_executor.Exit %7#1 : tensor<*x!tf_type.int32ref> {device = "", T = "tfdtype$DT_INT32"} loc("while/Exit")
    %10:2 = tf_executor.island(%7#2) wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_INT32", value = dense<1> : tensor<i32>} : () -> tensor<i32> loc("while/Add/y")

  }
  func.return %0#0, %0#1 : tensor<*xi32>, tensor<*xf32>
}

// Verify that While op is mapped to TensorFlow StatelessWhile op if the
// is_stateless attribute is present and otherwise it is mapped to TensorFlow
// While op. In both cases, the additional attribute should be dropped.

// CHECK: name: "StatefulWhile"
// CHECK-NOT: name:
// CHECK: op: "While"
// CHECK-NOT: is_stateless
// CHECK-NOT: shape_invariant
// CHECK:  attr {

// RUN: tf-opt --tfl-legalize-tf-while %s -o - | FileCheck %s
// RUN: tf-opt --tfl-legalize-tf-while %s -o - --tfl-legalize-tf-while --inline='default-pipeline=''' | FileCheck %s --check-prefix=INLINE
// RUN: tf-opt --tfl-legalize-tf-while %s -o - --tfl-legalize-tf-while --inline | FileCheck %s --check-prefix=CANON

  func.setPrivate();
}

void RunOnWhile(TF::WhileOp while_op) {
  Operation* op = while_op.getOperation();
  // Create new TFL While op that will be used to replace TF While op.
  auto new_op = OpBuilder(op).create<TFL::WhileOp>(
      op->getLoc(), op->getResultTypes(), op->getOperands(),
      while_op.getIsStateless());
  Location loc = while_op->getLoc();
  CreateRegionWithCall(while_op.cond_function(), new_op.getCond(), loc);

# Verify that TensorFlow While and StatelessWhile ops are mapped to the
# composite While op in MLIR with is_stateless attribute set accordingly to
# distinguish between them.

# CHECK-DAG: "tf.While"{{.*}} is_stateless = false{{.*}} loc(fused["While:", "StatefulWhile"])
# CHECK-DAG: "tf.While"{{.*}} is_stateless = true{{.*}} loc(fused["StatelessWhile:", "StatelessWhile"])

      while_op.getLoc(), new_result_types, new_operands, while_op->getAttrs());
  new_while_op.getCond().takeBody(while_op.getCond());
  new_while_op.getBody().takeBody(while_op.getBody());

  for (auto i = 0; i < n; ++i) {
    if (!while_op.getResult(i).use_empty()) {
      auto j = id_map[i];
      while_op.getResult(i).replaceAllUsesWith(new_while_op.getResult(j));
    }
  }

  auto new_while_op = OpBuilder(while_op).create<WhileOp>(
      while_op.getLoc(), new_types, operands, while_op->getAttrs());
  new_while_op.getCond().takeBody(while_op.getCond());
  new_while_op.getBody().takeBody(while_op.getBody());
  while_op.replaceAllUsesWith(
      new_while_op.getResults().take_front(while_op.getNumResults()));
  while_op.erase();
}

void WhileOutlinePass::runOnOperation() {

Dmitrii Gridin <******@****.***> 1716981664 +0200

}

// Rewrites the while op with extra chaining operands and results. Uses a
// dummy constant of requested type as argument to all the new chaining
// operands.
TF::WhileOp RewriteWhileOp(TF::WhileOp while_op, int num_resource_inputs,
                           ShapedType const_type) {
  IslandOp while_wrapper = while_op->getParentOfType<IslandOp>();
  assert(while_wrapper && "While op is expected to be wrapped in a IslandOp");