DatasetInput GetDatasetInput(Value value) {
  // TODO(haoliang): add an interface for DatasetOp to avoid the following
  // enumeration.
  // Iteratively tracing upwards if parent op is `IdentityOp` or `IdentityNOp`.
  while (
      llvm::isa_and_nonnull<IdentityOp, IdentityNOp>(value.getDefiningOp())) {
    value = value.getDefiningOp()->getOperand(
        mlir::cast<OpResult>(value).getResultNumber());
  }

          tf_executor.yield %identity : tensor<i32>
        }
        %island_1_result, %island_1_control = tf_executor.island {
          %identity_n:2 = "tf.IdentityN"(%arg1, %island_0_result) : (tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
          tf_executor.yield %identity_n#0
        }
        tf_executor.fetch %island_0_result, %island_1_result : tensor<i32>, tensor<i32>
      }

namespace mlir {
namespace TF {

namespace {
// Returns the equivalent Value skipping through identity nodes.
Value LookThroughIdentity(Value result) {
  while (isa_and_nonnull<IdentityOp, IdentityNOp>(result.getDefiningOp())) {
    auto op_result = result.cast<OpResult>();
    result = op_result.getOwner()->getOperand(op_result.getResultNumber());
  }
  return result;
}