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
\|
=> 3 (Max Pool)
=> 0 3 0 (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
|/ \|
=> 2 1 (Min Pool)
=> 2 0 1 (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
+ 2 0 1
-------
= 2 3 1
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 and the odd-even rules have to be swapped.
2 3 1 -> 2 0 3 (Max)
-> 1 (Min)
-----
2 1 3
Following up with a non-offset swap layer again completes the example,
2 1 3 -> 1 3
-> 2
-----
1 2 3
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 $nC2$ 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)$$
Linear activations will not work for negative numbers but you can use ReLU if you assume that a large positive number has been added to the input via the bias and then subtracted at the output with another layer.
Thus, sorting can be done with $O(n^2)$ layers with $n$ neurons each.