}

StringRef GetEntryFunctionName(TF::XlaCallModuleOp op) {
  if (!op->hasAttrOfType<FlatSymbolRefAttr>(
          TF::kStablehloEntryFunctionAttrName)) {
    return StringRef();
  }
  return op
      ->getAttrOfType<FlatSymbolRefAttr>(TF::kStablehloEntryFunctionAttrName)
      .getValue();
}

bool IsHybridQuantizedOp(Operation* op) {
  if ((op->getNumOperands() != 2 && op->getNumOperands() != 3) ||

// returns `std::nullopt`.
std::optional<llvm::StringRef> GetTpuClusterName(Operation* op) {
  if (auto tpu_status = op->getAttrOfType<StringAttr>(kTpuStatusAttr)) {
    // Borrow cluster name from TPU status (for `TPUCompilationResult` op).
    return tpu_status.getValue();
  }
  auto device_type = op->getAttrOfType<StringAttr>(TF::kCompileDeviceTypeAttr);
  if (!device_type || device_type.getValue() != TF::kTpuDevice) {

  for (Value v : live_ins) operand_types.emplace_back(v.getType());

  auto func_type = builder->getFunctionType(operand_types, op.getResultTypes());

  std::string func_name;
  if (auto outlined_func_name = op->template getAttrOfType<StringAttr>(
          TF::kClusterOutlinedFunctionNameAttr)) {
    op->removeAttr(TF::kClusterOutlinedFunctionNameAttr);
    func_name = outlined_func_name.str();
  } else {
    func_name = "_func";
  }

  for (Operation& op : *block) {
    // Collect device types which currently must be consistent per block
    // (checked later).
    if (auto device_type_attr =
            op.getAttrOfType<StringAttr>(mlir::TF::kCompileDeviceTypeAttr)) {
      // tf.StatefulPartitionedCall ops with and without
      // _tpu_replicate attributes may exist in the same graph. Ops without

}

void EncodeSharding(mlir::Operation* op, llvm::StringRef shard_str) {
  if (!op->hasAttrOfType<mlir::StringAttr>(shard_str)) return;

  ::xla::OpSharding sharding;
  auto sharding_proto_str =
      op->getAttrOfType<mlir::StringAttr>(shard_str).getValue().str();
  if (!sharding.ParseFromString(sharding_proto_str)) return;

  auto hlosharding = xla::HloSharding::FromProto(sharding);
  if (!hlosharding.ok()) {

    cost_map_[op] = iter->second(context, op);
    return;
  }

  // Try to use the recorded cost if any.
  if (cost_recorder_ != nullptr) {
    const auto op_key_attr =
        op->getAttrOfType<mlir::IntegerAttr>(kOpKeyAttrName);
    if (op_key_attr) {
      cost_map_[op] = cost_recorder_->GetCost(op_key_attr.getInt());
      return;
    }
  }

                                        tensor_cast.getResult());
  return success();
}

LogicalResult ConvertLSTMCellSimpleToFusedLSTM::InitializeFromFuncAttributes() {
  auto attr = fused_func_op_->getAttrOfType<StringAttr>(kTFImplements);
  if (!attr) {
    return fused_func_op_.emitError()
           << "Invalid function attribute, expected " << kTFImplements
           << " attribute "
              "not found";
  }

// Returns the `shared_name` attribute value if exists. If not, returns an
// empty string.
StringRef GetSharedName(Operation* op) {
  if (!op->hasAttrOfType<StringAttr>(kSharedNameAttr)) return "";
  return op->getAttrOfType<StringAttr>(kSharedNameAttr).getValue();
}

// Checks if the HashTable is initialized. This function assumes that the
// HashTable is initialized if it appears in the initializer since it can't

// here is to check if the graph is of the "main" function and does not have the
// "tf.entry_function" attribute defined.
bool CanPruneGraph(func::FuncOp func) {
  return func.getName() != "main" ||
         func->getAttrOfType<DictionaryAttr>("tf.entry_function") != nullptr;
}

// Visits an op's operand if it is an output of an Operation in the same
// tf_executor.graph.
void VisitOpOperand(GraphOp graph, Value operand,

};

// Check if op directly uses a key in `virtual_devices`.
bool DirectUseOfVirtualDevice(const DictionaryAttr& virtual_devices,
                              Operation* op) {
  StringAttr op_device = op->getAttrOfType<StringAttr>(kDeviceAttr);
  if (!op_device) return false;
  if (virtual_devices.get(op_device.getValue())) return true;
  return false;
}

// Check if op or its ancestor uses a key in `virtual_devices`.