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,

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

    auto chain_arg = func.getArgument(chain_index++);
    auto src_identity = builder_chain_src.create<TF::IdentityOp>(
        chain_arg.getLoc(), chain_arg.getType(), chain_arg);
    auto chain_src_island = CreateIsland(src_identity, {}, builder_chain_src);

    auto sink_identity = builder_chain_sink.create<TF::IdentityOp>(
        chain_arg.getLoc(), chain_arg.getType(), chain_arg);
    auto chain_sink_island =

  }
};

struct ConvertIdentity : public OpConversionPattern<TF::IdentityOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(
      TF::IdentityOp op, OpAdaptor adaptor,
      ConversionPatternRewriter &rewriter) const override {
    Value input = adaptor.getOperands()[0];
    rewriter.replaceOpWithNewOp<TF::IdentityOp>(
        op, input.getType(), adaptor.getOperands(), op->getAttrs());

// quantizable but is necessary to make the conversion successful.
bool IsAlwaysAllowlistedOp(Operation *op) {
  return llvm::isa<
      // clang-format off
      // go/keep-sorted start
      TF::ConstOp,
      TF::IdentityOp,
      TF::PartitionedCallOp,
      TF::StatefulPartitionedCallOp
      // go/keep-sorted end
      // clang-format on
      >(op);
}

// LINT.IfChange

// as GraphDef.
std::unique_ptr<OperationPass<ModuleOp>>
CreateMergeInitializerFunctionOpsToMainPass();

// Creates a pass that moves & merges the "@tf_quant__save" function to "@main"
// function. A new `IdentityOp` will be created. It will have control dependency
// to the save function and returns the file_prefix argument (typed
// `tensor<!tf_type.string>`). The file_prefix argument, which can be identified

// of Identity op is only used by TF ops.
struct RemoveIdentity : public OpRewritePattern<TF::IdentityOp> {
  using OpRewritePattern<TF::IdentityOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TF::IdentityOp identity,
                                PatternRewriter &rewriter) const override {
    // Replace the op with the input if input and result have the same type.

  func.return %1 : tensor<*xf32>
// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<1.000000e-03> : tensor<f32>}> {device = ""} : () -> tensor<f32>
// CHECK: %[[IDENTITY_0:.*]] = "tf.Identity"(%arg0) {device = ""} : (tensor<2xf32>) -> tensor<*xf32>

  func.return %0#1 : tensor<?xf32>
}

func.func @whileSinkConstant(%arg0: tensor<1x256xf32>) -> tensor<1x256xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "prefix", outputs = "Identity_1"}} {
    %cst_0 = arith.constant dense<1> : tensor<256x256xi8>

    // TODO(b/269195256#comment19) change the warning for Identity op to error
    // when issue with input graph is resolved. Possible issue with python layer
    // inserting Identity op incorrectly.
    if (isa<TF::IdentityOp>(op)) {
      op->emitWarning("TF/XLA TPU bridge input check: found invalid op. ")
          << op->getName() << " can't be both xla and non-xla";
      return true;
    }