Methods, systems, and apparatus, including computer programs encoded on computer storage media, for generating a schedule for a computation graph. One of the methods includes obtaining data representing an input computation graph; processing the data representing the input computation graph using a graph neural network to generate one or more instance-specific proposal distributions; and generating a schedule for the input computation graph by performing an optimization algorithm in accordance with the one or more instance-specific proposal distributions.