For example, original program:

    ```mlir
    module {
      func @while_body {
        %input = "tf.IteratorGetNext"(...) {device = "/CPU:0"}: -> tensor<2x224x224x3xf32>
        %device_launch = "tf_device.cluster_func"(%input,...) {func = @_func,...)
        return ...
      }

        // expected-remark@above {{ID: 1}}
        // expected-remark@above {{Successors: {2}}}
         %0 = "tf.IteratorGetNext"(%handle) {_class = ["loc:@GeneratorDataset_2"], device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource>) -> tensor<!tf_type.string>
        // expected-remark@above {{ID: 2}}

    const auto& element_type = element_types->list().type(idx);
    return ConvertToMlirTensorType(output_shape, element_type, &builder);
  };

  if (node.type_string() == "IteratorGetNext" ||
      node.type_string() == "IteratorGetNextSync" ||
      node.type_string() == "MultiDeviceIteratorGetNextFromShard")
    return type_from_array_attr("output_shapes", "output_types");