// clang-format off
          // go/keep-sorted start
          TF::AvgPoolOp,
          TF::ConcatOp,
          TF::ConcatV2Op,
          TF::ExpandDimsOp,
          TF::IdentityNOp,
          TF::IdentityOp,
          TF::MaxPoolOp,
          TF::PadV2Op,
          TF::RankOp,
          TF::ReshapeOp,
          TF::SelectOp,
          TF::SelectV2Op,
          TF::ShapeNOp,

// CHECK: %[[IDENTITY_1:.*]] = "tf.Identity"(%[[W_1]]) : (tensor<2x3x3x1024xi8>) -> tensor<2x3x3x1024xi8>
// CHECK: %[[DEQUANTIZED_1:.*]] = "tf.PartitionedCall"(%[[IDENTITY_1]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform__}> : (tensor<2x3x3x1024xi8>) -> tensor<2x3x3x1024xf32>

template <typename T>
std::tuple<T, llvm::SmallVector<mlir::TF::IdentityOp, 4>> GetSingleUserOfType(
    OpResult result) {
  llvm::SmallVector<mlir::TF::IdentityOp, 4> identity_ops;

  do {
    Operation* user = result.hasOneUse() ? *result.getUsers().begin() : nullptr;
    if (auto t = llvm::dyn_cast_or_null<T>(user)) {
      return std::make_tuple(t, identity_ops);
    } else if (auto identity =

    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()))

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>

#include "mlir/IR/Matchers.h"  // from @llvm-project
#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project

namespace mlir {
namespace TF {

class IdentityOp;
class IdentityNOp;

// Returns the RankedTensorType for the given operand. TensorFlow constant ops
// may have non-static shape because the shape is not propagated during constant

// CHECK-DAG:   %[[IDENTITY_0:.*]] = "tf.Identity"(%[[PARTITIONEDCALL_0]])
// CHECK-DAG:   %[[IDENTITY_1:.*]] = "tf.Identity"(%[[PARTITIONEDCALL_1]])
// CHECK:   return %[[IDENTITY_0]], %[[IDENTITY_1]] : tensor<1xf32>, tensor<1xf32>
// CHECK: }
}

// -----

  }
}

// Checks if an operation can have TPU devices propagated through.
bool IsSupportedOpToSetDevice(Operation& op) {
  return IsSupportedExecutorOp(op) ||
         isa<TF::IdentityOp, TF::IdentityNOp, TF::ShapeOp>(op);
}

// Finds nonconflicting TPU device for an operation from its operands. If an
// operand has no device or a non TPU device, or if there are conflicting

  auto is_valid_broadcasting = [](ShapedType operand_ty, ShapedType identity_ty,
                                  ShapedType result_ty) -> bool {
    // Scalar identity is broadcastable to any operand shape, we only need to
    // check that operand has the same shape as a result.
    bool scalar_identity = identity_ty.hasRank() && identity_ty.getRank() == 0;
    if (scalar_identity) return operand_ty == result_ty;

  %0 = "tf.Identity"(%arg0) {_XlaSharding = ""} : (tensor<i32>) -> tensor<i32>
  // CHECK: return %[[OUTPUT]]
  func.return %0 : tensor<i32>
}

// CHECK-LABEL: func @identity_n
// CHECK-SAME:    (%[[ARG0:.*]]: tensor<i32>, %[[ARG1:.*]]: tensor<f32>)
func.func @identity_n(%arg0: tensor<i32>, %arg1: tensor<f32>) -> (tensor<i32>, tensor<f32>) {
  // CHECK-NOT: "tf.IdentityN"