val_bcast.getType(), block_arg))) {
    return failure();
  }

  OpBuilder before_cluster_builder(cluster);
  IdentityOp assigned_id = before_cluster_builder.create<IdentityOp>(
      val_bcast.getLoc(), block_arg.getType(), block_arg);
  assigned_id->setAttr(kICIWeightDistributionMlirBridgeMarker,
                       before_cluster_builder.getBoolAttr(true));

                  AttrType &max_value) const {
    Value min = tf_op.getMin(), max = tf_op.getMax();
    if (auto min_id = min.getDefiningOp<TF::IdentityOp>()) {
      min = min_id.getInput();
    }
    if (auto max_id = max.getDefiningOp<TF::IdentityOp>()) {
      max = max_id.getInput();
    }

    if (!matchPattern(min, m_Constant(&min_value))) {
      return false;
    }

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

  // Try to retrieve lengths before building the attributes (should fail)
  EXPECT_EQ(-1, TFE_OpGetInputLength(identityOp, "input", status));
  CHECK_NE(TF_OK, TF_GetCode(status)) << TF_Message(status);
  EXPECT_EQ(-1, TFE_OpGetOutputLength(identityOp, "output", status));

    %identity_out, %control_identity = tf_executor.island wraps "tf.Identity"(%replicated_out) {device = ""} : (tensor<4x4xf32>) -> tensor<4x4xf32>
    %partitioned_out:2, %partitioned_control = tf_executor.island wraps "tf.TPUPartitionedOutput"(%identity_out) {partition_dim = 0 : i64} : (tensor<4x4xf32>) -> (tensor<2x4xf32>, tensor<2x4xf32>)
    tf_executor.fetch

#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace quant {

LogicalResult RemoveIdentity::matchAndRewrite(TF::IdentityOp identity,
                                              PatternRewriter &rewriter) const {
  for (Operation *user : identity->getUsers()) {
    // Replace the op with the input if output is only used by TF ops.

class RankEquals<string rank> : Constraint<CPred<
  "RankEquals($0, " # rank # ")">>;

def IsFusibleWithBias : Constraint<CPred<
  "IsFusibleWithBiasOp($0.getDefiningOp())">>;

// Folds TF IdentityOp with constant input.
def RemoveConstIdentityOp : Pat<
  (TF_IdentityOp (TF_ConstOp $input)),
  (TF_ConstOp $input)>;

// Standardizes the Max and Min ops by moving constant value to rhs. This will

    n = 2 : i32} {
     // CHECK:      %[[IDENTITY_OUT:.*]] = "tf.Identity"(%[[RI_0]])
     // CHECK:      %[[RESOURCE_OUT:.*]] = "tf_device.launch"()
     // CHECK-SAME: TPU_REPLICATED_CORE_0
     // CHECK-NEXT:   %[[READ_OUT:.*]] = "tf.ReadVariableOp"(%[[IDENTITY_OUT]])
     // CHECK-NEXT:   tf_device.return %[[READ_OUT]]

    const Block& func_block = func.front();
    // Walk all operations and nested regions to find accessed gradient sources
    // for function arguments.
    func.walk([&](Operation* op) {
      if (llvm::isa<TF::IdentityOp, TF::IdentityNOp>(op)) {
        op->replaceAllUsesWith(op->getOperands());
        return;
      }
      if (auto grad = llvm::dyn_cast<TF::TensorArrayGradV3Op>(op)) {

        decomposed_partitioned_call_callees) {
  for (auto& op : llvm::make_early_inc_range(block->getOperations())) {
    // TODO(yuanzx): Add a pass to remove identities in device computation.
    if (llvm::isa<TF::IdentityOp, TF::IdentityNOp, TF::StopGradientOp>(&op)) {
      op.replaceAllUsesWith(op.getOperands());
      op.erase();
    } else if (auto list = llvm::dyn_cast<TF::EmptyTensorListOp>(&op)) {

// such nodes to carry information.
void RemovePassthroughOp(Block &block) {
  for (auto &op : llvm::make_early_inc_range(block)) {
    if (isa<TF::IdentityOp, TF::IdentityNOp>(op) || IsCastOfResource(op)) {
      op.replaceAllUsesWith(op.getOperands());
      op.erase();
    }
  }
}

using LocalVarOp = std::variant<TF::VarHandleOp, TF::MlirLocalVarOp>;