// computation corresponding to i-th logical device. If the attribute does not
// exist, the all inputs are placed on logical core 0.
mlir::LogicalResult ExtractInputsForLogicalDevices(
    int num_cores_per_replica, mlir::tf_device::ClusterFuncOp cluster_func,
    mlir::OpBuilder* builder,
    llvm::SmallVectorImpl<llvm::SmallVector<mlir::Value, 4>>* input_list);