Automatic Memory Management for Compute Graphs

A method includes obtaining a compute graph for computing a first tensor, identifying in the graph a reduction operation in at least one dimension of the first tensor, locating, at the operation, a cut point that cuts the graph into first and second portions, and determining a plurality of slices of the first tensor. The method also includes backpropagating the cut point through the graph to define a plurality of first graph pieces for the first portion, each particular first graph piece representing a computation of a particular slice of the plurality of slices based on a particular portion of a plurality of portions of a second tensor. The method further includes defining one or more second graph pieces to combine outputs of the first graph pieces, and executing the first graph pieces and the second graph pieces to execute the first portion of the compute graph.