differentiable_sort.py

import numpy as np

def softmax(a,b):
    return np.log(np.exp(a)+np.exp(b))

def softmin(a,b):
    return -softmax(-a, -b)

def softrank(a, b):
    return softmin(a,b), softmax(a,b)


def nbits(n):
    return int(np.log(n) / np.log(2))


def bitonic_matrices(n):
    layers = nbits(n)
    matrices = []
    for layer in range(layers+1):
                        
        for sub in reversed(range(1, layer+1)):
            l = np.zeros((n//2, n))
            r = np.zeros((n//2, n))
            map_l = np.zeros((n, n//2))
            map_r = np.zeros((n, n//2))
            out = 0
            for i in range(0,n,2**sub):
                for j in range(2**(sub-1)):
                    ix = i + j
                    a, b = ix, ix+(2**(sub-1))
                    way = (ix >> layer) & 1
                    l[out, a] = 1
                    r[out, b] = 1
                    if way:                                            
                        a, b = b,a                                                            
                    map_l[a, out] = 1
                    map_r[b, out] = 1                                        
                    out += 1
            matrices.append((l, r, map_l, map_r))            
    return matrices
        
matrices = bitonic_matrices(16)        

def bisort(matrices, x):
    y = x
    for low, high, map_l, map_r in matrices:
        l, r = low @ y, high @ y                
        y = map_l @ np.minimum(l,r) + map_r @ np.maximum(l,r)
        
    return y

def diff_bisort(matrices, x):
    y = x
    for low, high, map_l, map_r in matrices:
        y1, y2 = softrank(low @ y, high @ y)        
        y = map_l @ y1 + map_r @ y2        
    return y

for i in range(10):
    print(bisort(matrices, np.random.randint(0,200,16)))
    
for i in range(10):
    print(diff_bisort(matrices, np.random.randint(0,200,16)))