}
}

// -----

// 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)>;

    for (const auto& user : op.getUsers()) {
      if (user->template hasTrait<OpTrait::IsTerminator>()) return true;
      if (auto next_user = dyn_cast_or_null<TF::IdentityOp>(user)) {
        return (*(next_user->getResult(0).getUsers().begin()))
            ->template hasTrait<OpTrait::IsTerminator>();
      }
    }
    return false;
  }

  rewriter.setInsertionPointAfter(op);
  auto const_op =
      rewriter.create<TF::ConstOp>(op.getLoc(), new_type, tensor_proto_attr);
  auto new_identity_op = rewriter.create<TF::IdentityOp>(
      op->getLoc(), const_op.getType(), const_op);
  return new_identity_op.getResult();
}

Operation* LogicsForUniformDequanization(PatternRewriter& rewriter,

  Operation* op_to_fold = FindOperationOfType<TF::MulOp>(test_func);
  SmallVector<Value> results = ConstantFoldOpIfPossible(op_to_fold);
  EXPECT_THAT(results, SizeIs(1));
  // No constant-folding since the IdentityOp has `TF_NoConstantFold` trait.
  auto mul_op = dyn_cast_or_null<TF::MulOp>(results[0].getDefiningOp());
  EXPECT_THAT(mul_op, NotNull());
  // Even though the preceding CastOp is foldable, it shouldn't be folded since

          next_values_to_visit.push_back(
              while_op.getBody().front().getArgument(operand_number));
          continue;
        }

        if (llvm::isa<mlir::TF::IdentityOp, mlir::TF::CastOp,
                      mlir::TF::ReadVariableOp>(owner)) {
          next_values_to_visit.push_back(use.getOwner()->getResult(0));
          continue;
        }