Use same API for `compute` and `build_graph`