Depth concatenation using a matrix computation unit

1. A method comprising:

receiving a request to process network inputs to a neural network using an integrated circuit that performs neural network computations in hardware using a matrix computation unit, the neural network comprising a depth concatenation neural network layer that specifies a concatenation of an input tensor having dimensions x 1 by y 1 by z 1 and an input tensor having dimensions x 1 by y 1 by z 2 along a depth dimension to generate an output tensor having dimensions x 1 by y 1 by (z 1 +z 2 ); and

generating instructions that, when executed by the integrated circuit, cause the integrated circuit to, during processing of a network input by the neural network, generate a layer output tensor that satisfies the specification of the depth concatenation neural network layer by performing operations comprising:

for each spatial location in a first input tensor to the depth concatenation layer and a second input tensor to the depth concatenation layer:

obtaining a first input depth vector for the spatial location in the first input tensor and a second input depth vector for the spatial location in the second input tensor, wherein each of the first input depth vector and the second input depth vector are made up of one or more respective fixed size chunk vectors;

identifying a first chunk vector from among the chunk vectors in the first input depth vector and the second input depth vector, wherein the first chunk vector is a first chunk vector according to an ordering of the chunk vectors to include one or more non-padded entries and one or more padded entries;

multiplying, using the matrix computation unit, a second chunk vector that follows the first chunk vector in the ordering by a first shift weight matrix for the depth concatenation layer to generate a shifted second chunk vector that has zeroes as its first (max−n) entries and a first n entries of the second chunk vector as its remaining entries, wherein n is the number of padded entries in the first chunk vector and max is the fixed size of the chunk vectors; and

adding, using the matrix computation unit, the shifted second chunk vector and a modified first chunk vector to generate a concatenated chunk vector that is part of the layer output, the modified first chunk vector having the non-padded entries of the first chunk vector and zeroes for the remaining entries of the modified first chunk vector,

wherein the first shift weight matrix for the depth concatenation layer is a max by max matrix.

2. The method of claim 1 , the operations further comprising: for each chunk vector before the first chunk vector in the ordering, moving the chunk vector through the matrix computation unit as part of the layer output.

3. The method of claim 1 , the operations further comprising: moving the modified first chunk vector to a set of output sum-in registers of the matrix computation unit; and

wherein adding the shifted second chunk vector and the modified first chunk vector comprises:

moving the shifted second chunk vector into the set of output sum-in registers of the matrix computation unit while the modified first chunk vector is stored in the set of output sum-in registers of the matrix computation unit.

4. The method of claim 3 , wherein moving the modified first chunk vector comprises:

multiplying the first chunk vector by a modified identity weight matrix for the depth concatenation layer using the matrix computation unit.

5. The method of claim 4 , further comprising: generating the modified identity weight matrix for the depth concatenation layer; and

storing the modified identity weight matrix for the depth concatenation layer in a memory accessible to the integrated circuit.

6. The method of claim 1 , further comprising: generating the first shift weight matrix for the depth concatenation layer; and

storing the first shift weight matrix for the depth concatenation layer in a memory accessible to the integrated circuit.

7. The method of claim 1 , wherein the first shift weight matrix has all entries be zero except for a diagonal row of ones starting at the first entry of the (max−n+1)-th column of the matrix.

8. A system comprising one or more computers and one or more storage devices storing first instructions that when executed by the one or more computers cause the one or more computers to perform first operations comprising:

receiving a request to process network inputs to a neural network using an integrated circuit that performs neural network computations in hardware using a matrix computation unit, the neural network comprising a depth concatenation neural network layer that specifies a concatenation of an input tensor having dimensions x 1 by y 1 by z 1 and an input tensor having dimensions x 1 by y 1 by z 2 along a depth dimension to generate an output tensor having dimensions x 1 by y 1 by (z 1 +z 2 ); and

generating instructions that, when executed by the integrated circuit, cause the integrated circuit to, during processing of a network input by the neural network, generate a layer output tensor that satisfies the specification of the depth concatenation neural network layer by performing second operations comprising:

for each spatial location in a first input tensor to the depth concatenation layer and a second input tensor to the depth concatenation layer:

obtaining a first input depth vector for the spatial location in the first input tensor and a second input depth vector for the spatial location in the second input tensor, wherein each of the first input depth vector and the second input depth vector are made up of one or more respective fixed size chunk vectors;

identifying a first chunk vector from among the chunk vectors in the first input depth vector and the second input depth vector, wherein the first chunk vector is a first chunk vector according to an ordering of the chunk vectors to include one or more non-padded entries and one or more padded entries;

multiplying, using the matrix computation unit, a second chunk vector that follows the first chunk vector in the ordering by a first shift weight matrix for the depth concatenation layer to generate a shifted second chunk vector that has zeroes as its first (max−n) entries and a first n entries of the second chunk vector as its remaining entries, wherein n is the number of padded entries in the first chunk vector and max is the fixed size of the chunk vectors; and

adding, using the matrix computation unit, the shifted second chunk vector and a modified first chunk vector to generate a concatenated chunk vector that is part of the layer output, the modified first chunk vector having the non-padded entries of the first chunk vector and zeroes for the remaining entries of the modified first chunk vector,

wherein the first shift weight matrix for the depth concatenation layer is a max by max matrix.

9. The system of claim 8 , the second operations further comprising: for each chunk vector before the first chunk vector in the ordering, moving the chunk vector through the matrix computation unit as part of the layer output.

10. The system of claim 8 , the second operations further comprising:

moving the modified first chunk vector to a set of output sum-in registers of the matrix computation unit; and

wherein adding the shifted second chunk vector and the modified first chunk vector comprises:

moving the shifted second chunk vector into the set of output sum-in registers of the matrix computation unit while the modified first chunk vector is stored in the set of output sum-in registers of the matrix computation unit.

11. The system of claim 10 , wherein moving the modified first chunk vector comprises:

multiplying the first chunk vector by a modified identity weight matrix for the depth concatenation layer using the matrix computation unit.

12. The system of claim 11 , the first operations further comprising:

generating the modified identity weight matrix for the depth concatenation layer; and

storing the modified identity weight matrix for the depth concatenation layer in a memory accessible to the integrated circuit.

13. The system of claim 8 , the first operations further comprising:

generating the first shift weight matrix for the depth concatenation layer; and

storing the first shift weight matrix for the depth concatenation layer in a memory accessible to the integrated circuit.

14. The system of claim 8 , wherein the first shift weight matrix has all entries be zero except for a diagonal row of ones starting at the first entry of the (max−n+1)-th column of the matrix.

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

receiving a request to process network inputs to a neural network using an integrated circuit that performs neural network computations in hardware using a matrix computation unit, the neural network comprising a depth concatenation neural network layer that specifies a concatenation of an input tensor having dimensions x 1 by y 1 by z 1 and an input tensor having dimensions x 1 by y 1 by z 2 along a depth dimension to generate an output tensor having dimensions x 1 by y 1 by (z 1 +z 2 ); and

generating instructions that, when executed by the integrated circuit, cause the integrated circuit to, during processing of a network input by the neural network, generate a layer output tensor that satisfies the specification of the depth concatenation neural network layer by performing second operations comprising:

for each spatial location in a first input tensor to the depth concatenation layer and a second input tensor to the depth concatenation layer:

obtaining a first input depth vector for the spatial location in the first input tensor and a second input depth vector for the spatial location in the second input tensor, wherein each of the first input depth vector and the second input depth vector are made up of one or more respective fixed size chunk vectors;

identifying a first chunk vector from among the chunk vectors in the first input depth vector and the second input depth vector, wherein the first chunk vector is a first chunk vector according to an ordering of the chunk vectors to include one or more non-padded entries and one or more padded entries;

multiplying, using the matrix computation unit, a second chunk vector that follows the first chunk vector in the ordering by a first shift weight matrix for the depth concatenation layer to generate a shifted second chunk vector that has zeroes as its first (max−n) entries and a first n entries of the second chunk vector as its remaining entries, wherein n is the number of padded entries in the first chunk vector and max is the fixed size of the chunk vectors; and

adding, using the matrix computation unit, the shifted second chunk vector and a modified first chunk vector to generate a concatenated chunk vector that is part of the layer output, the modified first chunk vector having the non-padded entries of the first chunk vector and zeroes for the remaining entries of the modified first chunk vector,

wherein the first shift weight matrix for the depth concatenation layer is a max by max matrix.

16. The computer storage media of claim 15 , the second operations further comprising:

for each chunk vector before the first chunk vector in the ordering, moving the chunk vector through the matrix computation unit as part of the layer output.

17. The computer storage media of claim 15 , the second operations further comprising:

moving the modified first chunk vector to a set of output sum-in registers of the matrix computation unit; and

wherein adding the shifted second chunk vector and the modified first chunk vector comprises:

moving the shifted second chunk vector into the set of output sum-in registers of the matrix computation unit while the modified first chunk vector is stored in the set of output sum-in registers of the matrix computation unit.

18. The computer storage media of claim 17 , wherein moving the modified first chunk vector comprises:

multiplying the first chunk vector by a modified identity weight matrix for the depth concatenation layer using the matrix computation unit.

19. The computer storage media of claim 15 , the first operations further comprising:

generating the first shift weight matrix for the depth concatenation layer; and

storing the first shift weight matrix for the depth concatenation layer in a memory accessible to the integrated circuit.

20. The computer storage media of claim 15 , wherein the first shift weight matrix has all entries be zero except for a diagonal row of ones starting at the first entry of the (max−n+1)-th column of the matrix.