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 performing the requested operation; 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; the method further comprising: analyzing the computational graph to identify groups of nodes that operate on shared data flowing on directed edges to the group of nodes; wherein the partitioning comprises generating, for each identified group, a respective subgraph including the identified group of nodes; wherein a respective device to which the operations of the respective subgraph are assigned re-uses memory storing the same shared data for operations performed by nodes of the identified group of nodes of the respective subgraph. 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 performing the requested operation; 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; the method further comprising: analyzing the computational graph to identify groups of nodes that operate on shared data flowing on directed edges to the group of nodes; wherein the partitioning comprises generating, for each identified group, a respective subgraph including the identified group of nodes; wherein a respective device to which the operations of the respective subgraph are assigned re-uses memory storing the same shared data for operations performed by nodes of the identified group of nodes of the respective subgraph. analyzing the computational graph to identify groups of nodes that operate on shared data flowing on directed edges to the group of nodes; wherein the partitioning comprises generating, for each identified group, a respective subgraph including the identified group of nodes; wherein a respective device to which the operations of the respective subgraph are assigned re-uses memory storing the same shared data for operations performed by nodes of the identified group of nodes of the respective subgraph.