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

  TFE_OpAddInput(op, b, status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TF_DeleteStatus(status);
  TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(a));

  return op;
}

TFE_Op* IdentityOp(TFE_Context* ctx, TFE_TensorHandle* a) {
  TF_Status* status = TF_NewStatus();

  TFE_Op* op = TFE_NewOp(ctx, "Identity", status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  for (auto res : compilation_result) {
    // Build identity op with the same location/name as the original compilation
    // result op.
    result_id = builder->create<TF::IdentityOp>(
        res.getLoc(), compile_op->getResult(0).getType(),
        result_id->getResult(0));
    // Assign to same device as result is currently set, unless unset and then

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

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

}

// Trace the graph to find out the actual operation.
Value getActualValue(Operation *op) {
  if (isa<TF::CastOp>(op)) {
    op = op->getOperand(0).getDefiningOp();
  }

  if (isa<TF::IdentityOp>(op)) {
    op = op->getOperand(0).getDefiningOp();
  }
  return op->getResult(0);
}

Value CreateXlaDotV2OpFromTfEinsumOp(OpBuilder &builder, Location loc,