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

  while (auto result = mlir::dyn_cast<OpResult>(v)) {
    if (!(allow_other_use || v.hasOneUse())) break;
    auto op = result.getDefiningOp();
    if (!llvm::isa<TF::IdentityOp, TF::IdentityNOp>(op)) {
      break;
    }
    v = op->getOperand(result.getResultNumber());
    if (skipped) skipped->insert(op);
  }
  return v;
}

    TypeID::get<TF::_FusedConv2DOp>(),
    TypeID::get<TF::GatherNdOp>(),
    TypeID::get<TF::GatherV2Op>(),
    TypeID::get<TF::GreaterEqualOp>(),
    TypeID::get<TF::IdentityOp>(),
    TypeID::get<TF::IdentityNOp>(),
    TypeID::get<TF::InplaceUpdateOp>(),
    TypeID::get<TF::InvertPermutationOp>(),
    TypeID::get<TF::IRFFTOp>(),
    TypeID::get<TF::L2LossOp>(),
    TypeID::get<TF::LegacyCallOp>(),

        return ownerData != null
            ? identityFor(ownerData, relativeUrlTo(componentData.coordinates, ownerData.coordinates))
            : identityFor(componentData, null);
    }

    private ModuleMetadataSpec.Identity identityFor(ComponentData componentData, String relativeUrl) {
        return new ModuleMetadataSpec.Identity(

  // isn't a composite op. The following ops are explicitly skipped here because
  // their "no-op" expansion is known to cause problems in some cases.
  static const char* kOpsToSkip[] = {
      "IdentityOp",
      "NoOp",              // b/174596063
      "OptionalHasValue",  // b/173136483
      "OptionalGetValue",  // b/173136483
      "VarHandleOp",       // b/176819198
  };
  for (const char* skip : kOpsToSkip) {

# CHECK: return %[[IDENTITY]]#1, %[[IDENTITY]]#0


def Test():

  x = tf.constant(1.0, shape=(5, 3))
  y = tf.constant(1.0, shape=(3, 5))

  s = tf.matmul(x, y)
  t = tf.matmul(y, x)
  [t, s] = array_ops.identity_n([t, s])

  tensor_info_x = tf.compat.v1.saved_model.utils.build_tensor_info(x)
  tensor_info_y = tf.compat.v1.saved_model.utils.build_tensor_info(y)
  tensor_info_s = tf.compat.v1.saved_model.utils.build_tensor_info(s)

  return value.getDefiningOp();
}

// TODO(b/158596585): Replace this with a cost model analysis.
bool IsTrivialUnaryOperation(Operation* op) {
  return llvm::isa<TF::CastOp, TF::IdentityOp>(op);
}

// Adds outside compilation attributes to unary ops such as Identity/Cast ops
// at the head of TPU computation that is used only by other outside compiled

  %val1 = "tf.Identity"(%val0) : (tensor<96x!tf_type.f32ref>) -> tensor<96xf32>
  func.return
}

// -----

// Test case: Basic converting. '_class' attribute is at IdentityOp.

func.func @f() {
  // CHECK: "tf.VarHandleOp"
  // CHECK: "tf.ReadVariableOp"
  %val0 = "tf.VariableV2"() {container = "", device = "", shape = #tf_type.shape<96>, shared_name = ""} : () -> tensor<96x!tf_type.f32ref>