namespace mlir {
namespace TFTPU {

namespace {

// This pass removes Identity/IdentityN ops from the TPU computation and
// reachable functions.
// TODO(lyandy): Remove this pass once resource op lifting is migrated to use
// resource alias analysis and support region based control flow. Removing
// Identity ops may remove `_XlaSharding` annotation attribute if Identity ops
// are used to propagate such information.

  %3 = "tf.IdentityN"(%2) {device = ""} : (tensor<1x2xf32>) -> tensor<1x2xf32>
  return %3 : tensor<1x2xf32>
}

// CHECK: func @cast_bf16_matmul_to_fp32
// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<10x2xf32>}> : () -> tensor<10x2xf32>
// CHECK: %[[matmul:.*]] = "tf.MatMul"(%arg0, %[[cst]])
// CHECK: %[[identity:.*]] = "tf.IdentityN"(%[[matmul]])

// CHECK: %[[conv:.*]] = "tf.Conv2D"(%[[ARG0]], %[[cast]])
// CHECK: %[[identity:.*]] = "tf.IdentityN"(%[[conv]]) {device = ""} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32>
// CHECK: return %[[identity]] : tensor<1x3x2x2xf32>

// -----

// Tests that `tf.BatchFunction` is inlined.

jiuker <******@****.***> 1684487613 +0800

  func.return
}
// CHECK-LABEL: empty_island_multiple_data_results
// CHECK-SAME: (%[[ARG_0:.*]]: tensor<*xf32>, %[[ARG_1:.*]]: tensor<*xi32>)
// CHECK: %[[IDENTITY_N:.*]]:2 = "tf.IdentityN"
// CHECK-SAME: (%[[ARG_0]], %[[ARG_1]])
// CHECK: tf_executor.yield %[[IDENTITY_N]]#0, %[[IDENTITY_N]]#1

func.func @single_op_island_forward_block_arg(%arg0: tensor<?x?x?x?xbf16>) -> (tensor<2048xf32>, tensor<?x?x?x?xbf16>) {
  %0:2 = tf_executor.graph {

    return absl::OkStatus();
  }

  // Create a single IdentityN that is dead if and only if
  // `predicate_as_control` is dead.
  //
  // IdentityN is also special in that, unlike `Identity`, it does not place
  // int32 inputs in host memory.  Placing int32 inputs in host memory would
  // defeat the purpose of adding this indirection.
  ops::IdentityN identity_n(root.WithOpName("int32_id_n"), int32_inputs);

    auto identity = builder.create<TF::IdentityOp>(island.getLoc(),
                                                   operand.getType(), operand);
    yield.setOperand(0, identity.getOutput());
  } else {
    auto identity_n = builder.create<TF::IdentityNOp>(
        island.getLoc(), yield.getOperandTypes(), yield.getOperands());
    for (auto it : llvm::enumerate(identity_n.getResults()))

// By removing identity ops, constant operands with dynamic shapes have static
// shape information which is necessary for correct pattern matching in this
// pass.
struct RemoveIdentity : public OpRewritePattern<TF::IdentityOp> {
  using OpRewritePattern<TF::IdentityOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TF::IdentityOp identity,

    // H2D transfer kernels. So fold the tf.Identity op if:
    // * the identity op is placed on TPU, and
    // * the arg to the identity op is produced by a TPUExecuteOp.
    if (op_device && op_device.getValue().contains("TPU")) {
      return true;
    }

    Operation *def_op = arg.getDefiningOp();
    // If the arg to this identity op is the arg of a function, there's no
    // defining op.

func.func @while_main(%arg0: tensor<?x256x256xf32>) -> (tensor<i32>, tensor<256x256xf32>, tensor<?x256x256xf32>) attributes {tf.entry_function = {inputs = "input", outputs = "Identity,Identity_1,Identity_2"}} {
  %cst = arith.constant dense<1.000000e+00> : tensor<256x256xf32>
  %cst_0 = arith.constant dense<0> : tensor<i32>
  %cst_1 = arith.constant dense<-1> : tensor<i32>