It's possible to encode a version of Bubble Sort by hand, that can be shown to correctly sort numbers.
Bubble Sort proceeds by flipping adjacent elements of the array which are inverted. For example,
3 2 1
x
2 3 1
x
2 1 3
x
1 2 3
This can be implemented with a double for loop
for i in 0..n {
for j in i+1..n {
if arr[i] > arr[j] {
arr[i], arr[j] = arr[j], arr[i];
}
}
}
Swap Layers
To begin, we will design a swap layer which can swap adjacent elements to correct their order.
Let $X_i$ be the input vector encoded as floats. The swap layer $Y_i$ has the same size as the input. All equations are written with $0$ based indexing.
$$ Y_i = \begin{cases} min(X_i, X_{i+1}) & \text{if $i$ is even} \\ max(X_i, X_{i-1}) & \text{if $i$ is odd} \\ \end{cases} $$
The odd case needs a stride 2 max pool followed by transposed convolution that pads a zero on the left.
3 2 1 5
\| \|
=> 3 5 (Max Pool)
=> 0 3 0 5 (Padding)
The even case is just a stride 2 min pool (explained later) followed by a transposed convolution layer that pads the input with 0 on the right,
3 2 1 5
|/ |/
=> 2 1 (Min Pool)
=> 2 0 1 0 (Zero padding)
The two pools are then summed to produce a swap layer. Summing can be done with a simple 1x1 conv layer with stride 1.
0 3 0 5
+ 2 0 1 0
-------
= 2 3 1 5
The overall structure of a swap layer looks like this,
Input --> MinPool -- Zero pad-- + --> Swapped
| |
|----> MaxPool -- Zero pad --|
Notice the swap layer works locally in a 1x2 receptive field. To allow bubble sort like movement across fields, the next swap layer has be shifted by 1 position, thus inverting the odd-even rule above.
2 3 1 5 -> 0 3 0 5 (Max)
-> 2 0 1 0 (Min)
-------
2 3 1 5
Following up with a non-offset swap layer again completes the example,
2 1 3 5 -> 1 0 3 0 (Min)
-> 0 2 0 5 (Max)
-------
1 2 3 5
Stacking enough swap layers would eventually sort the array since each pair of swap layers will fix at least one inversion and there are at most $^{n}C_2$ inversions.
The min pool can be implemented with 2 1x1 convolution layers with linear activations and 1 max pool in between since,
$$min(a, b) = -max(-a, -b)$$
ReLU activations can't pass negative numbers but you can use ReLU if you assume that a large positive number has been added to the inputs to make them all positive and then subtracted out at the output.
Thus, sorting can be done with $O(n^2)$ layers with $n$ neurons each.