A computer-implemented method of performing a neural network or machine learning operation, comprising: receiving a request from a client to process a computational graph representing a neural network or machine learning model; obtaining data representing the computational graph, the computational graph comprising a plurality of nodes and directed edges representing the neural network or machine learning model, wherein each node represents a respective operation, wherein each directed edge connects a respective first node to a respective second node that represents an operation that receives, as input, an output of an operation represented by the respective first node; identifying a plurality of available devices for processing the computational graph, wherein the available devices comprise devices having a processing unit and memory and include devices having different amounts of memory; partitioning the computational graph into a plurality of subgraphs, each subgraph comprising one or more nodes in the computational graph; and assigning, for each subgraph, the operations represented by the one or more nodes in the subgraph to a respective available device in the plurality of available devices for processing, wherein the assigning to the respective available device includes determining the device to which the subgraph is assigned by estimating a maximum amount of memory to be used by any node in the subgraph; wherein estimating the maximum amount of memory to be used by any node in the subgraph comprises traversing the subgraph to calculate a dimension of a tensor on each directed edge to and from each node of the subgraph; wherein the assigning comprises assigning the subgraph to a device which has memory capable of storing a largest tensor flowing in the subgraph as indicated by the dimension of the tensor; and wherein the request identifies an output from a respective node in the computational graph, and wherein the method further comprising: for each of the devices to which a subgraph was assigned, causing the device to perform the operations represented by the nodes of the subgraph assigned to the device; and after a final node of the subgraph assigned to the device performs operations, providing the output to the client. receiving a request from a client to process a computational graph representing a neural network or machine learning model; obtaining data representing the computational graph, the computational graph comprising a plurality of nodes and directed edges representing the neural network or machine learning model, wherein each node represents a respective operation, wherein each directed edge connects a respective first node to a respective second node that represents an operation that receives, as input, an output of an operation represented by the respective first node; identifying a plurality of available devices for processing the computational graph, wherein the available devices comprise devices having a processing unit and memory and include devices having different amounts of memory; partitioning the computational graph into a plurality of subgraphs, each subgraph comprising one or more nodes in the computational graph; and assigning, for each subgraph, the operations represented by the one or more nodes in the subgraph to a respective available device in the plurality of available devices for processing, wherein the assigning to the respective available device includes determining the device to which the subgraph is assigned by estimating a maximum amount of memory to be used by any node in the subgraph; wherein estimating the maximum amount of memory to be used by any node in the subgraph comprises traversing the subgraph to calculate a dimension of a tensor on each directed edge to and from each node of the subgraph; wherein the assigning comprises assigning the subgraph to a device which has memory capable of storing a largest tensor flowing in the subgraph as indicated by the dimension of the tensor; and wherein the request identifies an output from a respective node in the computational graph, and wherein the method further comprising: for each of the devices to which a subgraph was assigned, causing the device to perform the operations represented by the nodes of the subgraph assigned to the device; and after a final node of the subgraph assigned to the device performs operations, providing the output to the client. for each of the devices to which a subgraph was assigned, causing the device to perform the operations represented by the nodes of the subgraph assigned to the device; and after a final node of the subgraph assigned to the device performs operations, providing the output to the client.