for (OpType op : func.getOps<OpType>()) {
    if (found_op != nullptr) {
      func.emitOpError() << "number of " << found_op.getOperationName()
                         << " in loop body is not 1";
      return LogicalResult::failure();
    }
    if (GetReplicationAttr(op) != replication_attr) {
      op.emitOpError() << "is not part of the replication region "

        if (func_arg.getOwner() != &block)
          return read_op.emitOpError(kResourceFunctionMsg);

        ResourceInfo& resource_info = resources[func_arg.getArgNumber()];
        resource_info.read = true;
        read_op.getValue().replaceAllUsesWith(resource_info.live_value);
      } else {
        return read_op.emitOpError(kInvalidResourceMsg);
      }

      read_op.erase();

  if (!num_cores_per_replica_attr)
    return cluster.emitOpError(
        CreateMissingAttributeMsg(tensorflow::kNumCoresPerReplicaAttr));
  int num_cores_per_replica = num_cores_per_replica_attr.getInt();

  auto topology_attr = cluster->getAttrOfType<StringAttr>("topology");
  if (!topology_attr)
    return cluster.emitOpError(
        CreateMissingAttributeMsg(tensorflow::kTopologyAttr));

      if (arg_it == info.stack_var_arg_to_size_arg.end()) continue;
      auto it = data_var_to_size_var.find(call.getOperand(i));
      if (it == data_var_to_size_var.end()) {
        call.emitOpError("unknown stack");
        return failure();
      }
      assert(arg_it->second == new_operands.size());
      new_operands.push_back(it->getSecond());
    }
    OpBuilder builder(call);

  auto embedding_attr = op->getAttrOfType<StringAttr>(kTPUEmbeddingAttr);
  if (!embedding_attr) return mlir::success();

  if (!embedding_op_map->insert({embedding_attr.getValue(), op}).second)
    return op->emitOpError(
        "found duplicate TPU embedding ops potentially from multiple "
        "TPUEmbedding layers");

  op->removeAttr(kTPUEmbeddingAttr);
  return success();
}

LogicalResult FindTPUEmbeddingOps(

      } else if (device_str_attr.getValue() != launch.getDevice()) {
        return launch.emitOpError()
               << "inner op has conflicting 'device' attribute, "
                  "got '"
               << device_str_attr.getValue() << "' but expected '"
               << launch.getDevice() << "'";
      }
    } else {
      return launch.emitOpError()
             << "inner op has bad 'device' attribute, got " << device_attr;
    }

        element_type) {
      return op->emitOpError(
          "requires compatible element types for all operands and results");
    }
  }
  // Verify that all operand element types are compatible to `element_type`.
  for (const auto& operand_type : op->getOperandTypes()) {
    if (mlir::tf_type::GetElementTypeOrSelfResolveRef(operand_type) !=
        element_type) {
      return op->emitOpError(

  Block* parent_block = graph.getOperation()->getBlock();
  for (Operation& op : graph.GetBody().without_terminator()) {
    auto island_op = llvm::dyn_cast<tf_executor::IslandOp>(op);
    if (!island_op)
      return op.emitOpError()
             << "is not supported for lifting out of tf_executor.graph, "
                "expected tf_executor.island";

    // Move inner ops in island to before the outer graph.

      GetPartitionedOutputsAndIdentityOps(cluster);

  for (auto [_, ops] : partitioned_outputs) {
    if (!(ops.empty() || ops.size() == num_replicas)) {
      return (ops.begin())->emitOpError()
             << "expected zero or " << num_replicas
             << " 'TPUPartitionedOutput' op(s), instead got "
             << partitioned_outputs.size();
    }
  }

  // No need to replicate.

        base_dilations_vec, window_dilations_vec);
    if (!window) {
      op->emitOpError("failed to create window");
    }
    auto output_shape = InferWindowOutputShape(
        input_ty, window.value(),
        mlir::cast<ShapedType>(op.getInitValue().getType()).getElementType());

    if (!output_shape) {
      op->emitOpError("failed to infer output shape");
    }