// CHECK-DAG: %[[ITER0:.*]]:2 = "tf.IteratorGetNext"(%[[ARG0]])
  // CHECK-DAG: %[[ITER1:.*]]:2 = "tf.IteratorGetNext"(%[[ARG1]])
  %2:2 = "tf.IteratorGetNext"(%arg0) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> (tensor<128xf32>, tensor<128xf32>)
  %3:2 = "tf.IteratorGetNext"(%arg1) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> (tensor<128xf32>, tensor<128xf32>)

  op: "_Arg"
  attr {
    key: "T"
    value {
      type: DT_INT64
    }
  }
  attr {
    key: "index"
    value {
      i: 4
    }
  }
}
node {
  name: "IteratorGetNext"
  op: "IteratorGetNext"
  input: "args_0"
  attr {
    key: "output_shapes"
    value {
      list {
        shape {
          dim {
            size: 1
          }
          dim {
            size: 8

*   At last, the rewrite pass will insert SpaceToDepthOp after IteratorGetNext
    where the iterator's result has the same shape as the first convolution's
    input.

#### Pseudo MLIR code before and after RewritePass

```mlir
// Example: original program:
//
module {
   func @while_body {
     %input = "tf.IteratorGetNext"(...) {device = "/CPU:0"}:
              -> tensor<2x224x224x3xf32>

// artificial dependencies and incur great performance loss.  In this example,
// Node_Y becomes dependent on IteratorGetNext and the latencies add up if
// Node_X and Node_Y are in the same cluster.
//
// IteratorGetNext -> Node_X -> Stage
//
// Unstage -> Node_Y
//
Status ClusterScopingPassImpl::ScopingForPipelineStages() {
  for (Node* n : graph_->nodes()) {

    %0 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    // CHECK: %[[INPUT:.*]] = "tf.IteratorGetNext"
    %1 = "tf.IteratorGetNext"(%arg5) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource>) -> tensor<2x224x224x3xf32>

  %2:2 = "tf.IteratorGetNext"(%0) {_class = ["loc:@iterator"], device = ""} : (tensor<!tf_type.resource>) -> (tensor<200x28x28x1xf32>, tensor<200x10xf32>)
  // CHECK: tf.Iterator
  // CHECK: tf.MakeIterator
  // CHECK: tf.IteratorGetNext
  return

  func.func private @get_next(%arg0: tensor<!tf_type.resource>) -> tensor<200x10xf32> {
    // CHECK: %0 = "tf.Iterator"
    // CHECK-SAME: shared_name = "foo_iterator"
    // CHECK: "tf.IteratorGetNext"(%0)
    %0 = "tf.IteratorGetNext"(%arg0) : (tensor<!tf_type.resource>) -> tensor<200x10xf32>
    return %0 : tensor<200x10xf32>
  }

  // CHECK-LABEL: @main
  func.func @main()

  REGISTER_KERNEL_BUILDER(Name("DeleteIterator").Device(DEVICE),               \
                          data::DeleteIteratorOp);                             \
  REGISTER_KERNEL_BUILDER(Name("IteratorGetNext").Device(DEVICE),              \
                          data::IteratorGetNextOp);                            \
  REGISTER_KERNEL_BUILDER(Name("IteratorGetNextAsOptional").Device(DEVICE),    \

 Cast 38
 ConcatV2 3
 Const 875
 ControlTrigger 5
 Enter 874
 Equal 4
 Exit 69
 ExpandDims 9
 Fill 5
 FloorMod 1
 GreaterEqual 7
 Identity 113
 IsVariableInitialized 1
 IteratorGetNext 1
 IteratorV2 1
 Less 9
 LogicalAnd 3
 LoopCond 8
 Max 4
 Maximum 44
 Merge 145
 Minimum 43
 Mul 8
 NextIteration 136
 RandomUniform 14
 Range 9
 RefSwitch 166
 Reshape 2

    return "Inserts TPU layout ops to determine layout at run time.";
  }
};

// Checks if the input producer op is supported in this transform. Right now, we
// only check if it is a tf.IteratorGetNext where resource input is coming from
// a VarHandle on CPU or a function argument assigned to CPU.
bool IsSupportedInputOp(
    Operation* op,
    const TF::ResourceAliasAnalysis::Info& resource_alias_analysis) {