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());
  }

    for (auto resource_user : resource_users_to_visit) {
      add_resource_op_to_list(resource_user);

      // Account for pass-through identity ops.
      if (auto pass_through_identity =
              llvm::dyn_cast<TF::IdentityOp>(resource_user)) {
        for (auto identity_user :
             pass_through_identity.getOutput().getUsers()) {
          new_resource_users.emplace_back(identity_user);
        }
      }
    }

    llvm::StringMap<PartitionedCallStackOpsInfo>*
        decomposed_partitioned_call_callees) {
  for (auto& op : llvm::make_early_inc_range(block->getOperations())) {
    if (llvm::isa<TF::IdentityOp, TF::IdentityNOp>(&op)) {
      // Removes identity nodes in the block. The device computation does not
      // need such nodes to carry information.
      op.replaceAllUsesWith(op.getOperands());
      op.erase();

  func.return
}
// CHECK-LABEL: empty_island_multiple_data_results
// CHECK-SAME: (%[[ARG_0:.*]]: tensor<*xf32>, %[[ARG_1:.*]]: tensor<*xi32>)
// CHECK: %[[IDENTITY_N:.*]]:2 = "tf.IdentityN"
// CHECK-SAME: (%[[ARG_0]], %[[ARG_1]])
// CHECK: tf_executor.yield %[[IDENTITY_N]]#0, %[[IDENTITY_N]]#1

func.func @single_op_island_forward_block_arg(%arg0: tensor<?x?x?x?xbf16>) -> (tensor<2048xf32>, tensor<?x?x?x?xbf16>) {
  %0:2 = tf_executor.graph {

  }
}

// -----

// Test that the @tf_quant_save's ops are cloned to @main. When there are no
// __tf_file_prefix argument in @main, confirm that it is created and wired
// to the newly created `IdentityOp`.

module attributes {tf_saved_model.semantics} {
  func.func private @tf_quant__save(%arg: tensor<!tf_type.string>) -> () {
    tf_executor.graph {

// Return a matmul op multiplying `a` by `b`.
TFE_Op* MatMulOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b);

// Return an identity op.
TFE_Op* IdentityOp(TFE_Context* ctx, TFE_TensorHandle* a);

// Return a shape op fetching the shape of `a`.
TFE_Op* ShapeOp(TFE_Context* ctx, TFE_TensorHandle* a);

    Value value = nullptr;
    for (auto v : op.getComponents()) {
      value = v;
    }
    if (!value) return failure();
    rewriter.replaceOpWithNewOp<TF::IdentityOp>(op, value.getType(), value);
    return success();
  }
};

class HandleOptionalGet : public OpRewritePattern<TF::OptionalGetValueOp> {
 public:
  using OpRewritePattern::OpRewritePattern;

  // defeat the purpose of adding this indirection.
  ops::IdentityN identity_n(root.WithOpName("int32_id_n"), int32_inputs);
  root.graph()->AddControlEdge(predicate_as_control.node(),
                               identity_n.operation.node());

  for (int i = 0, end = int32_inputs.size(); i < end; i++) {
    TF_RETURN_IF_ERROR(root.graph()->UpdateEdge(identity_n[i].node(), i, n,

func.func @empty_island_multiple_data_results(%arg0: tensor<*xf32>, %arg1: tensor<*xi32>) {
  tf_executor.graph {
    %0:3 = tf_executor.island {
      // CHECK: %[[IDENTITY_N:.*]]:2 = "tf.IdentityN"
      // CHECK-SAME: (%[[ARG_0]], %[[ARG_1]])
      // CHECK: tf_executor.yield %[[IDENTITY_N]]#0, %[[IDENTITY_N]]#1
      tf_executor.yield %arg0, %arg1 : tensor<*xf32>, tensor<*xi32>
    }
    tf_executor.fetch
  }
  func.return
}

  [(IsReshapeEquivalentToTranspose $transpose_out1, $output),
   (AnyStaticShapeTensor $input),
   (AnyStaticShapeTensor $output),
   (HasOneUse $transpose_out1)]>;

// Partially supported in TFLite, treated as passthrough IdentityOp
def ConvertCheckNumerics : Pat<(TF_CheckNumericsOp $arg, $msg), (TF_IdentityOp $arg)>;
def ConvertSnapshot : Pat<(TF_SnapshotOp $arg), (TF_IdentityOp $arg)>;