class IdentityOp;
class IdentityNOp;

// Returns if a value corresponds to a constant, returns the matched constant
// as an attribute.
template <typename AttrT>
bool GetValueAsConstant(Value val, AttrT &attr) {
  while (auto result = mlir::dyn_cast<OpResult>(val)) {
    Operation *op = result.getOwner();
    if (!isa<IdentityOp>(op) && !isa<IdentityNOp>(op)) break;

      // Control output is generated by the IslandOp, not the yield in
      // in the Island body.
      if (value == island.getControl()) break;
      value = island.GetYield().getOperand(res_index);
    } else if (isa<IdentityNOp, IdentityOp>(op)) {
      value = op->getOperand(res_index);
    } else if (auto call = dyn_cast<CallOpInterface>(op)) {
      func::FuncOp func = dyn_cast<func::FuncOp>(

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

  Operation* op = value.getDefiningOp();

    }
    return success();
  }
};

// Bypasses IdentityN op.
class ConvertIdentityNOp : public OpRewritePattern<TF::IdentityNOp> {
 public:
  using OpRewritePattern<TF::IdentityNOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(TF::IdentityNOp op,
                                PatternRewriter &rewriter) const override {
    rewriter.replaceOp(op, op.getOperands());
    return success();