Distributing tensor computations across computing devices

1. A computer-implemented method comprising:

receiving specification data that specifies a distribution of tensor computations among a plurality of computing devices,

wherein each tensor computation (i) is defined as receiving, as input, one or more respective input tensors each having one or more respective input dimensions, (ii) is defined as generating, as output, one or more respective output tensors each having one or more respective output dimensions, or (iii) is defined as both receiving, as input, one or more respective input tensors each having one or more respective input dimensions and generating, as output, one or more respective output tensors each having one or more respective output dimensions, and

wherein the specification data specifies a respective layout for each input and output tensor that assigns each dimension of the input or output tensor to one or more of the plurality of computing devices;

assigning, based on the layouts for the input and output tensors, respective device-local operations to each of the plurality of computing devices, comprising determining that the layout, if assigned, does not cause data to be lost for a respective tensor corresponding to the layout; and

causing the tensor computations to be executed by the plurality of computing devices by causing each of the plurality of computing devices to execute at least the respective-device local operations assigned to the computing devices.

2. The method of claim 1 , wherein the tensor computations are operations for training a machine learning model on training data using a machine learning training technique, and

wherein the specification data specifies respective layouts for each tensor processed while performing an iteration of the machine learning training technique on a batch of training data.

3. The method of claim 1 , wherein the tensor computations are represented by a directed graph of nodes connected by directed edges, each node representing a respective tensor computation and each incoming edge to a node representing an input tensor to the tensor computation represented by the node and each outgoing edge from the node representing an output tensor of the tensor computation represented by the node, and

wherein the specification data specifies respective layouts for each of the tensors represented by edges in the directed graph.

4. The method of claim 1 , wherein, for each input and output tensor, the layout specifies, for each dimension of the tensor, either that (i) the data along the dimension is replicated across all of the computing devices in the plurality of computing devices or that (ii) the data along the dimension is split among at least two of the plurality of computing devices.

5. The method of claim 4 , wherein the plurality of computing devices are arranged into an n-dimensional array, and wherein, for each input and output tensor, the layout specifies, for each dimension of the tensor, either that (i) the data along the dimension is replicated across all of the computing devices in the plurality of devices or that (ii) the data along the dimension is split between the computing devices along one of the dimensions of the n-dimensional array.

6. The method of claim 5 , wherein a shape of each of input and output tensor is referred to as a tuple of shape symbols, and wherein the specification data identifies, for at least one of the shape symbols, a mapping from a shape symbol to an identifier for one of n dimensions.

7. The method of claim 1 , wherein assigning respective device-local operations to the plurality of computing devices further comprises, for each tensor computation:

identifying, for each input tensor to the tensor computation and for each device, a slice of the input tensor that is assigned to the device by the specification data; and

assigning, to each device, a slice-wise instance of the tensor computation that applies the tensor computation to the slices of the input tensors that are assigned to the device.

8. The method of claim 1 , wherein assigning respective device-local operations to the plurality of computing devices comprises:

determining, for a first tensor computation, that communication between devices is required for slice-wise instances of the first tensor computation to generate a correct output for the first tensor computation, and

assigning, to the plurality of devices, one or more communication operations that cause communication between devices to combine outputs generated by the slice-wise instances of the first tensor computation.

9. The method of claim 8 , wherein assigning communication operations that cause communication between devices to combine outputs generated by the slice-wise instances of the operations comprises inserting an Allreduce operation.

10. The method of claim 1 , wherein the specification data is specified by a user, and wherein the method further comprises:

determining that a first layout of the layouts, if assigned, causes data to be lost for a respective tensor corresponding to the first layout, and in response, notifying the user that the first layout is an illegal layout.

11. 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 specification data that specifies a distribution of tensor computations among a plurality of computing devices,

wherein each tensor computation (i) is defined as receiving, as input, one or more respective input tensors each having one or more respective input dimensions, (ii) is defined as generating, as output, one or more respective output tensors each having one or more respective output dimensions, or (iii) is defined as both receiving, as input, one or more respective input tensors each having one or more respective input dimensions and generating, as output, one or more respective output tensors each having one or more respective output dimensions; and

wherein the specification data specifies a respective layout for each input and output tensor that assigns each dimension of the input or output tensor to one or more of the plurality of computing devices;

assigning, based on the layouts for the input and output tensors, respective device-local operations to each of the plurality of computing devices, comprising determining that the layout, if assigned, does not cause data to be lost for a respective tensor corresponding to the layout; and

causing the tensor computations to be executed by the plurality of computing devices by causing each of the plurality of computing devices to execute at least the respective-device local operations assigned to the computing devices.

12. The system of claim 11 , wherein the tensor computations are operations for training a machine learning model on training data using a machine learning training technique, and

wherein the specification data specifies respective layouts for each tensor processed while performing an iteration of the machine learning training technique on a batch of training data.

13. The system of claim 11 , wherein the tensor computations are represented by a directed graph of nodes connected by directed edges, each node representing a respective tensor computation and each incoming edge to a node representing an input tensor to the tensor computation represented by the node and each outgoing edge from the node representing an output tensor of the tensor computation represented by the node, and

wherein the specification data specifies respective layouts for each of the tensors represented by edges in the directed graph.

14. The system of claim 11 , wherein, for each input and output tensor, the layout specifies, for each dimension of the tensor, either that (i) the data along the dimension is replicated across all of the computing devices in the plurality of computing devices or that (ii) the data along the dimension is split among at least two of the plurality of computing devices.

15. The system of claim 14 , wherein the plurality of computing devices are arranged into an n-dimensional array, and wherein, for each input and output tensor, the layout specifies, for each dimension of the tensor, either that (i) the data along the dimension is replicated across all of the computing devices in the plurality of devices or that (ii) the data along the dimension is split between the computing devices along one of the dimensions of the n-dimensional array.

16. The system of claim 15 , wherein a shape of each of input and output tensor is referred to as a tuple of shape symbols, and wherein the specification data identifies, for at least one of the shape symbols, a mapping from a shape symbol to an identifier for one of n dimensions.

17. The system of claim 11 , wherein assigning respective device-local operations to the plurality of computing devices further comprises, for each tensor computation:

identifying, for each input tensor to the tensor computation and for each device, a slice of the input tensor that is assigned to the device by the specification data; and

assigning, to each device, a slice-wise instance of the tensor computation that applies the tensor computation to the slices of the input tensors that are assigned to the device.

18. The system of claim 11 , wherein assigning, based on the assignments of the dimensions of the input and output tensors, respective device-local operations to the plurality of computing devices comprises:

determining, for a first tensor computation, that communication between devices is required for slice-wise instances of the first tensor computation to generate a correct output for the first tensor computation, and

assigning, to the plurality of devices, one or more communication operations that cause communication between devices to combine outputs generated by the slice-wise instances of the first tensor computation.

19. The system of claim 18 , wherein assigning communication operations that cause communication between devices to combine outputs generated by the slice-wise instances of the operations comprises inserting an Allreduce operation.

20. One or more non-transitory computer-readable storage media encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform operations comprising:

receiving specification data that specifies a distribution of tensor computations among a plurality of computing devices,

wherein each tensor computation (i) is defined as receiving, as input, one or more respective input tensors each having one or more respective input dimensions, (ii) is defined as generating, as output, one or more respective output tensors each having one or more respective output dimensions, or (iii) is defined as both receiving, as input, one or more respective input tensors each having one or more respective input dimensions and generating, as output, one or more respective output tensors each having one or more respective output dimensions, and

wherein the specification data specifies a respective layout for each input and output tensor that assigns each dimension of the input or output tensor to one or more of the plurality of computing devices;

assigning, based on the layouts for the input and output tensors, respective device-local operations to each of the plurality of computing devices, comprising determining that the layout, if assigned, does not cause data to be lost for a respective tensor corresponding to the layout; and

causing the tensor computations to be executed by the plurality of computing devices by causing each of the plurality of computing devices to execute at least the respective-device local operations assigned to the computing devices.