Generating and processing augmented computational graphs

A system comprising:

• one or more computers and one or more storage devices on which are stored instructions that are operable, when executed by the one or more computers, to cause the one or more computers to perform operations comprising:
• receiving a request to process a first computational graph, the first computational graph comprising
• (i) a plurality of nodes and directed edges, each node representing a respective operation of a neural network modeled by the first computational graph and each directed edge connecting a respective first node to a respective second node, the respective first node for each directed edge representing a respective first operation of the neural network that produces data for processing according to a respective second operation of the neural network represented by the respective second node;
• and
• (ii) a reference to a function indicating that an output of a first particular node in the first computational graph is to be provided as input to the function and that an output of the function is to be provided as input to a second particular node in the first computational graph, wherein the function represents a second computational graph comprising a second plurality of nodes and directed edges, wherein the second computational graph is a portion of a third computational graph that has been designated as the function;
• in response to receiving the request to process the first computational graph:
• (i) generating an augmented computational graph that integrates the second computational graph represented by the function into the first computational graph;
• and
• (ii) processing the augmented computational graph including assigning different portions of the augmented computational graph to a plurality of processing devices;
• and
• wherein the function is a re-usable function that allows for the second computational graph to be integrated multiple times without regenerating the second computational graph each time.

One or more non-transitory computer-readable media on which are stored instructions that are operable, when executed by one or more computers, to cause the one or more computers to perform operations comprising:

• receiving a request to process a first computational graph, the first computational graph comprising
• (i) a plurality of nodes and directed edges, each node representing a respective operation of a neural network modeled by the first computational graph and each directed edge connecting a respective first node to a respective second node, the respective first node for each directed edge representing a respective first operation of the neural network that produces data for processing according to a respective second operation of the neural network represented by the respective second node;
• and
• (ii) a reference to a function indicating that an output of a first particular node in the first computational graph is to be provided as input to the function and that an output of the function is to be provided as input to a second particular node in the first computational graph, wherein the function represents a second computational graph comprising a second plurality of nodes and directed edges, wherein the second computational graph is a portion of a third computational graph that has been designated as the function;
• in response to receiving the request to process the first computational graph:
• (i) generating an augmented computational graph that integrates the second computational graph represented by the function into the first computational graph;
• and
• (ii) processing the augmented computational graph including assigning different portions of the augmented computational graph to a plurality of processing devices;
• and
• wherein the function is a re-usable function that allows for the second computational graph to be integrated multiple times without regenerating the second computational graph each time.