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>

    // 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.

    if (!all_of(variable_v2_op.getResult().getUsers(), [&variable_v2_op](
                                                           Operation *user) {
          if (!isa<IdentityOp>(user)) {
            variable_v2_op.emitOpError()
                << "expects all users to be 'tf.Identity', but got user "
                << user->getName();
            return false;
          }
          return true;
        })) {

  // CHECK-NOT: tf.Identity
  %5 = "tf.Identity"(%4) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "/device:CPU:0"} : (tensor<3x3xf32>) -> tensor<3x3xf32>
  // CHECK-NOT: tf.IdentityN

Aditya Manthramurthy <******@****.***> 1714599073 -0700

//
// input   min cst       max cst          input   min cst       max cst
//  \       |             |                \       |             |
//   \  (tf.Identity) (tf.Identity)   =>    \  (tf.Identity) (tf.Identity)
//    \     |             |                  \     |             |
//       tf.FakeQuantWithMinMaxVars       tf.FakeQuantWithMinMaxVars
//                   |                                 |

    call_op_returns.append(call_op.getResults().begin(),
                           call_op.getResults().end());
    SetFunctionPrivate(func_op);
  }

  // Creates Identity/IdentityN ops for returing values. This allows us to
  // restore the same output tensor names in python.
  int32_t output_count = 0;
  // Map from node name to the list of the OutputInfos of its outputs that are

    ```mlir
    func @my_fn(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
      %identity = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>
      %identity_n:2 = "tf.IdentityN"(%arg1, %identity) : (tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
      return %identity, %identity_n#0 : tensor<i32>, tensor<i32>
    }
    ```
  }];

// convert the output type to the next op. Here are the transformations:
//
// input   min cst       max cst              input
//  \       |             |                     |
//   \  (tf.Identity) (tf.Identity)   =>   quant.qcast
//    \     |             |                     |
//       tf.FakeQuantWithMinMaxVars        quant.dcast
//                   |                          |
//

                        visited_wrapped_ops, /*incoming=*/false)) {
    return failure();
  }
  return success();
}

// Whenever we find an Identity op that is unqualified, we remove this Identity
// op from the list `tpu_ops`. An unqualified Identity op indicates either its
// inputs or its outputs do not belong to the same cluster.
bool ExcludeIdentityOp(llvm::SmallDenseSet<Operation*>& tpu_ops,